Update auth.py

auth.py

@@ -1,10 +1,6 @@
 import os
 import uuid
 import logging
-import secrets
-import smtplib
-from email.mime.multipart import MIMEMultipart
-from email.mime.text import MIMEText
 from datetime import datetime, timedelta
 from urllib.parse import quote_plus
 from typing import List, Optional, Any
@@ -19,11 +15,16 @@ from pydantic import BaseModel, EmailStr, Field, validator
 from pymongo import MongoClient
 import os.path
 
-# Optionally import markdown to convert markdown text to HTML
+# ----- OTP and Email Imports -----
+import smtplib
+from email.mime.multipart import MIMEMultipart
+from email.mime.text import MIMEText
 try:
     import markdown
 except ImportError:
     markdown = None
+import secrets
+# ----------------------------------
 
 load_dotenv()
 
@@ -37,6 +38,8 @@ MONGO_URL = os.getenv("CONNECTION_STRING").replace("${PASSWORD}", password)
 client = MongoClient(MONGO_URL)
 db = client.users_database
 users_collection = db.users
+# New collection to store OTP records
+otp_collection = db.otp_verifications
 
 # OAuth2 setup
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
@@ -44,7 +47,7 @@ oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 # Create an APIRouter instance
 router = APIRouter()
 
-# ------------------ Pydantic Models ------------------
+# Pydantic models
 class User(BaseModel):
     name: str = Field(..., min_length=3, max_length=50)
     email: EmailStr
@@ -96,12 +99,7 @@ class LoginResponse(Token):
 class TokenData(BaseModel):
     email: Optional[str] = None
 
-# Model for OTP-based login
-class OTPLogin(BaseModel):
-    email: EmailStr
-    otp: str
-
-# ------------------ Password Hashing ------------------
+# Password hashing
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
 def verify_password(plain_password: str, hashed_password: str) -> bool:
@@ -110,7 +108,6 @@ def verify_password(plain_password: str, hashed_password: str) -> bool:
 def get_password_hash(password: str) -> str:
     return pwd_context.hash(password)
 
-# ------------------ User & Auth Helpers ------------------
 def get_user(email: str) -> Optional[dict]:
     return users_collection.find_one({"email": email})
 
@@ -148,7 +145,7 @@ def get_current_user(token: str = Depends(oauth2_scheme)) -> dict:
     except JWTError:
         raise HTTPException(status_code=401, detail="Invalid token")
 
-# ------------------ Avatar File Saving ------------------
+# Setup for avatar file saving
 AVATAR_DIR = "avatars"
 if not os.path.exists(AVATAR_DIR):
     os.makedirs(AVATAR_DIR)
@@ -176,7 +173,7 @@ def save_avatar_file(file: UploadFile) -> str:
         file.file.close()
     return file_path
 
-# ------------------ OTP and Email Functions ------------------
+# ----- OTP and Email Functions -----
 def generate_otp(length=6):
     """Generate a numeric OTP of specified length."""
     digits = "0123456789"
@@ -234,9 +231,44 @@ def send_email_func(sender_email, sender_password, receiver_email, subject, body
     except Exception as e:
         logger.error(f"Error sending email: {e}")
         raise e
-# -------------------------------------------------------------
+# ------------------------------------
 
-# ------------------ Auth Endpoints ------------------
+# ----- Auth Endpoints -----
+
+@router.post("/send_otp")
+async def send_otp(receiver_email: EmailStr = Body(..., embed=True)):
+    """
+    Generates an OTP and sends it via email to the provided receiver_email.
+    (For demonstration purposes, the OTP is returned. Remove it in production.)
+    """
+    sender_email = os.getenv("SENDER_EMAIL")
+    sender_password = os.getenv("SENDER_PASSWORD")
+    if not sender_email or not sender_password:
+        raise HTTPException(status_code=500, detail="Email sender credentials are not configured.")
+    
+    otp = generate_otp(6)
+    subject = "Your One-Time Password (OTP)"
+    body = f"""\
+# Your OTP
+
+Here is your one-time password (OTP):
+
+`{otp}`
+
+Please use the code above to verify your email and proceed with account creation.
+"""
+    try:
+        send_email_func(sender_email, sender_password, receiver_email, subject, body)
+        # Store OTP record with a 10-minute expiration
+        otp_record = {
+            "email": receiver_email,
+            "otp": otp,
+            "expires_at": datetime.utcnow() + timedelta(minutes=1)
+        }
+        otp_collection.update_one({"email": receiver_email}, {"$set": otp_record}, upsert=True)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error sending OTP email: {str(e)}")
+    return {"message": "OTP sent successfully", "otp": otp}  # Remove OTP from response in production
 
 @router.post("/signup", response_model=Token)
 async def signup(
@@ -244,8 +276,20 @@ async def signup(
     name: str = Form(...),
     email: EmailStr = Form(...),
     password: str = Form(...),
+    otp: str = Form(...),
     avatar: Optional[UploadFile] = File(None)
 ):
+    # Verify OTP for this email
+    otp_record = otp_collection.find_one({"email": email})
+    if not otp_record:
+        raise HTTPException(status_code=400, detail="No OTP sent to this email. Please request an OTP first.")
+    if otp_record["otp"] != otp:
+        raise HTTPException(status_code=400, detail="Invalid OTP provided.")
+    if datetime.utcnow() > otp_record["expires_at"]:
+        raise HTTPException(status_code=400, detail="OTP has expired. Please request a new one.")
+    # Remove the OTP record after successful verification
+    otp_collection.delete_one({"email": email})
+    
     try:
         _ = User(name=name, email=email, password=password)
     except Exception as e:
@@ -272,87 +316,21 @@ async def signup(
         "token_type": "bearer"
     }
 
-@router.post("/login_otp", response_model=LoginResponse)
-async def login_otp(otp_data: OTPLogin):
-    """
-    OTP-based login endpoint.
-    The user must supply their email and the OTP received via email.
-    """
-    user = get_user(otp_data.email)
+@router.post("/login", response_model=LoginResponse)
+async def login(request: Request, form_data: OAuth2PasswordRequestForm = Depends()):
+    user = authenticate_user(form_data.username, form_data.password)
     if not user:
-        logger.warning(f"Login OTP attempt for non-existent email: {otp_data.email}")
-        raise HTTPException(status_code=401, detail="User not found")
-    
-    # Check if OTP exists and is valid
-    stored_otp = user.get("otp")
-    otp_expires = user.get("otp_expires")
-    if not stored_otp or not otp_expires:
-        raise HTTPException(status_code=400, detail="OTP not generated. Please request a new OTP.")
-    
-    # Check if OTP is expired (using ISO format stored in DB)
-    if datetime.utcnow() > datetime.fromisoformat(otp_expires):
-        raise HTTPException(status_code=400, detail="OTP has expired. Please request a new OTP.")
-    
-    if otp_data.otp != stored_otp:
-        raise HTTPException(status_code=401, detail="Invalid OTP.")
-    
-    # OTP is valid. Remove OTP fields.
-    users_collection.update_one(
-        {"email": otp_data.email},
-        {"$unset": {"otp": "", "otp_expires": ""}}
-    )
-    logger.info(f"User logged in with OTP: {otp_data.email}")
+        logger.warning(f"Failed login attempt for: {form_data.username}")
+        raise HTTPException(status_code=401, detail="Incorrect username or password")
+    logger.info(f"User logged in: {user['email']}")
     return {
-        "access_token": create_access_token(otp_data.email),
-        "refresh_token": create_refresh_token(otp_data.email),
+        "access_token": create_access_token(user["email"]),
+        "refresh_token": create_refresh_token(user["email"]),
         "token_type": "bearer",
         "name": user["name"],
         "avatar": user.get("avatar")
     }
 
-@router.post("/send_otp")
-async def send_otp(receiver_email: EmailStr = Body(..., embed=True)):
-    """
-    Generates an OTP, stores it in the user's document with an expiration, 
-    and sends it via email to the provided receiver_email.
-    (For demonstration purposes, the OTP is returned. Remove in production.)
-    """
-    user = get_user(receiver_email)
-    if not user:
-        raise HTTPException(status_code=404, detail="User not found. Please sign up first.")
-    
-    sender_email = os.getenv("SENDER_EMAIL")
-    sender_password = os.getenv("SENDER_PASSWORD")
-    if not sender_email or not sender_password:
-        raise HTTPException(status_code=500, detail="Email sender credentials are not configured.")
-    
-    otp = generate_otp(6)
-    # Set OTP to expire in 5 minutes
-    otp_expires = (datetime.utcnow() + timedelta(minutes=1)).isoformat()
-    subject = "Your One-Time Password (OTP)"
-    body = f"""\
-# Your OTP
-
-Here is your one-time password (OTP):
-
-`{otp}`
-
-Please use the code above to proceed with your login.
-"""
-    try:
-        send_email_func(sender_email, sender_password, receiver_email, subject, body)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error sending OTP email: {str(e)}")
-    
-    # Update user document with OTP and expiration
-    users_collection.update_one(
-        {"email": receiver_email},
-        {"$set": {"otp": otp, "otp_expires": otp_expires}}
-    )
-    
-    # For demonstration, return the OTP (remove in production)
-    return {"message": "OTP sent successfully", "otp": otp}
-
 @router.get("/user/data")
 async def get_user_data(request: Request, current_user: dict = Depends(get_current_user)):
     return {