Update auth.py

auth.py

@@ -1,10 +1,13 @@
+# auth.py
 import os
 import uuid
 import logging
 from datetime import datetime, timedelta
+from urllib.parse import quote_plus
 from typing import Optional
 
-from fastapi import APIRouter, HTTPException, Request, UploadFile, File, Form, Query, Depends
+from dotenv import load_dotenv
+from fastapi import APIRouter, HTTPException, Depends, Request, UploadFile, File, Form
 from fastapi.responses import StreamingResponse
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
 from jose import JWTError, jwt
@@ -15,14 +18,16 @@ import gridfs
 from models import User, UserUpdate, Token, LoginResponse
 from config import CONNECTION_STRING, SECRET_KEY, ACCESS_TOKEN_EXPIRE_MINUTES, REFRESH_TOKEN_EXPIRE_DAYS
 
-# Configure logger
+load_dotenv()
+
 logger = logging.getLogger("uvicorn")
 logger.setLevel(logging.INFO)
 
-# Initialize MongoDB
+# Updated MongoDB initialization: now using CONNECTION_STRING from config.py
 client = MongoClient(CONNECTION_STRING)
 db = client.users_database
 users_collection = db.users
+# GridFS instance for storing avatars
 fs = gridfs.GridFS(db, collection="avatars")
 
 # OAuth2 setup
@@ -51,7 +56,8 @@ def create_token(data: dict, expires_delta: timedelta = None) -> str:
     to_encode = data.copy()
     expire = datetime.utcnow() + (expires_delta or timedelta(minutes=15))
     to_encode.update({"exp": expire})
-    return jwt.encode(to_encode, SECRET_KEY, algorithm="HS256")
+    algorithm = "HS256"
+    return jwt.encode(to_encode, SECRET_KEY, algorithm=algorithm)
 
 def create_access_token(email: str) -> str:
     return create_token({"sub": email}, timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES))
@@ -60,22 +66,16 @@ def create_refresh_token(email: str) -> str:
     return create_token({"sub": email}, timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS))
 
 def get_current_user(token: str = Depends(oauth2_scheme)) -> dict:
-    logger.info(f"[AUTH] Received token: {token}")
     try:
         payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
-        logger.info(f"[AUTH] Decoded payload: {payload}")
         email: str = payload.get("sub")
         if not email:
-            logger.warning("[AUTH] Token payload missing 'sub'")
             raise HTTPException(status_code=401, detail="Invalid credentials")
         user = get_user(email)
         if not user:
-            logger.warning(f"[AUTH] No user found for email: {email}")
             raise HTTPException(status_code=401, detail="User not found")
-        logger.info(f"[AUTH] Authenticated user: {email}")
         return user
-    except JWTError as e:
-        logger.error(f"[AUTH] JWT decode failed: {e}")
+    except JWTError:
         raise HTTPException(status_code=401, detail="Invalid token")
 
 async def save_avatar_file_to_gridfs(file: UploadFile) -> str:
@@ -86,10 +86,14 @@ async def save_avatar_file_to_gridfs(file: UploadFile) -> str:
             status_code=400,
             detail="Invalid image format. Only JPEG, PNG, and GIF are accepted."
         )
-    contents = await file.read()
-    file_id = fs.put(contents, filename=file.filename, contentType=file.content_type)
-    logger.info(f"Avatar stored in GridFS with file_id: {file_id}")
-    return str(file_id)
+    try:
+        contents = await file.read()
+        file_id = fs.put(contents, filename=file.filename, contentType=file.content_type)
+        logger.info(f"Avatar stored in GridFS with file_id: {file_id}")
+        return str(file_id)
+    except Exception as e:
+        logger.exception("Failed to store avatar in GridFS")
+        raise HTTPException(status_code=500, detail="Could not store avatar file in MongoDB.")
 
 @router.post("/signup", response_model=Token)
 async def signup(
@@ -99,7 +103,6 @@ async def signup(
     password: str = Form(...),
     avatar: Optional[UploadFile] = File(None)
 ):
-    # Validate and create user
     try:
         _ = User(name=name, email=email, password=password)
     except Exception as e:
@@ -109,12 +112,22 @@ async def signup(
         logger.warning(f"Attempt to register already existing email: {email}")
         raise HTTPException(status_code=400, detail="Email already registered")
     hashed_password = get_password_hash(password)
-    user_data = {"name": name, "email": email, "hashed_password": hashed_password, "chat_histories": []}
+    user_data = {
+        "name": name,
+        "email": email,
+        "hashed_password": hashed_password,
+        "chat_histories": []
+    }
     if avatar:
-        user_data["avatar"] = await save_avatar_file_to_gridfs(avatar)
+        file_id = await save_avatar_file_to_gridfs(avatar)
+        user_data["avatar"] = file_id
     users_collection.insert_one(user_data)
     logger.info(f"New user registered: {email}")
-    return {"access_token": create_access_token(email), "refresh_token": create_refresh_token(email), "token_type": "bearer"}
+    return {
+        "access_token": create_access_token(email),
+        "refresh_token": create_refresh_token(email),
+        "token_type": "bearer"
+    }
 
 @router.post("/login", response_model=LoginResponse)
 async def login(request: Request, form_data: OAuth2PasswordRequestForm = Depends()):
@@ -122,37 +135,28 @@ async def login(request: Request, form_data: OAuth2PasswordRequestForm = Depends
     if not user:
         logger.warning(f"Failed login attempt for: {form_data.username}")
         raise HTTPException(status_code=401, detail="Incorrect username or password")
-    avatar_url = f"/auth/avatar/{user['avatar']}" if user.get('avatar') else None
     logger.info(f"User logged in: {user['email']}")
-    return {"access_token": create_access_token(user['email']), "refresh_token": create_refresh_token(user['email']), "token_type": "bearer", "name": user['name'], "avatar": avatar_url}
-
-@router.get("/user/data", response_model=User)
-async def get_user_data(
-    request: Request,
-    email: Optional[str] = Query(None, description="User email")
-):
-    # Determine if Authorization header is provided
-    auth_header = request.headers.get("Authorization", "")
-    user = None
-    if auth_header.startswith("Bearer "):
-        token = auth_header.split(" ", 1)[1]
-        # Authenticated path
-        user = get_current_user(token)
-    elif email:
-        # Public path
-        user = get_user(email)
-        if not user:
-            raise HTTPException(status_code=404, detail="User not found")
-    else:
-        raise HTTPException(status_code=422, detail="Provide either Authorization header or 'email' query parameter")
+    avatar_url = None
+    if "avatar" in user and user["avatar"]:
+        avatar_url = f"/auth/avatar/{user['avatar']}"
+    return {
+        "access_token": create_access_token(user["email"]),
+        "refresh_token": create_refresh_token(user["email"]),
+        "token_type": "bearer",
+        "name": user["name"],
+        "avatar": avatar_url
+    }
 
-    # Build response
-    avatar_url = f"/auth/avatar/{user['avatar']}" if user.get('avatar') else None
+@router.get("/user/data")
+async def get_user_data(request: Request, current_user: dict = Depends(get_current_user)):
+    avatar_url = None
+    if "avatar" in current_user and current_user["avatar"]:
+        avatar_url = f"/auth/avatar/{current_user['avatar']}"
     return {
-        "name": user['name'],
-        "email": user['email'],
+        "name": current_user["name"],
+        "email": current_user["email"],
         "avatar": avatar_url,
-        "chat_histories": user.get('chat_histories', [])
+        "chat_histories": current_user.get("chat_histories", [])
     }
 
 @router.put("/user/update")
@@ -165,23 +169,24 @@ async def update_user(
     current_user: dict = Depends(get_current_user)
 ):
     update_data = {}
-    if name:
-        update_data['name'] = name
-    if email:
-        update_data['email'] = email
-    if password:
+    if name is not None:
+        update_data["name"] = name
+    if email is not None:
+        update_data["email"] = email
+    if password is not None:
         try:
-            _ = User(name=current_user['name'], email=current_user['email'], password=password)
+            _ = User(name=current_user["name"], email=current_user["email"], password=password)
         except Exception as e:
             logger.error(f"Password validation error during update: {e}")
             raise HTTPException(status_code=400, detail=str(e))
-        update_data['hashed_password'] = get_password_hash(password)
+        update_data["hashed_password"] = get_password_hash(password)
     if avatar:
-        update_data['avatar'] = await save_avatar_file_to_gridfs(avatar)
+        file_id = await save_avatar_file_to_gridfs(avatar)
+        update_data["avatar"] = file_id
     if not update_data:
         logger.info("No update parameters provided")
         raise HTTPException(status_code=400, detail="No update parameters provided")
-    users_collection.update_one({'email': current_user['email']}, {'$set': update_data})
+    users_collection.update_one({"email": current_user["email"]}, {"$set": update_data})
     logger.info(f"User updated: {current_user['email']}")
     return {"message": "User updated successfully"}
 
@@ -195,6 +200,7 @@ from bson import ObjectId
 @router.get("/avatar/{file_id}")
 async def get_avatar(file_id: str):
     try:
+        # Convert the file_id string to an ObjectId before fetching
         file = fs.get(ObjectId(file_id))
         return StreamingResponse(file, media_type=file.content_type)
     except Exception as e: