Update core/security.py

core/security.py

@@ -1,61 +1,68 @@
 from datetime import datetime, timedelta
 from passlib.context import CryptContext
 from jose import jwt, JWTError
-from fastapi import Depends, HTTPException, status
+from fastapi import Depends, HTTPException, status, Request
 from fastapi.security import OAuth2PasswordBearer
 from core.config import SECRET_KEY, ALGORITHM, ACCESS_TOKEN_EXPIRE_MINUTES
 from db.mongo import users_collection
 import logging
-from fastapi import Request
 
 logger = logging.getLogger(__name__)
 
 # OAuth2 setup
 oauth2_scheme = OAuth2PasswordBearer(
-    tokenUrl="/auth/login",  # Correct path
+    tokenUrl="/auth/login",
     scheme_name="JWT"
 )
+
 # Password hashing context
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
-# Hash a plain password
 def hash_password(password: str) -> str:
     return pwd_context.hash(password)
 
-# Verify a plain password against the hash
 def verify_password(plain: str, hashed: str) -> bool:
     return pwd_context.verify(plain, hashed)
 
-# Create a JWT access token
-def create_access_token(data: dict, expires_delta: timedelta = None):
+def create_access_token(data: dict, expires_delta: timedelta = None) -> str:
     to_encode = data.copy()
     expire = datetime.utcnow() + (expires_delta or timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES))
     to_encode.update({"exp": expire})
     return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
 
-# Get the current user from the JWT token
 async def get_current_user(request: Request, token: str = Depends(oauth2_scheme)):
-    print("🧪 Request headers:", dict(request.headers))
-    print("🔐 Raw token received:", token)
+    credentials_exception = HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Could not validate credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+    
+    try:
+        # First try to get token from Authorization header
+        if not token:
+            # Fallback to checking cookies if using session-based auth
+            token = request.cookies.get("access_token")
+            if not token:
+                logger.error("No token provided in headers or cookies")
+                raise credentials_exception
 
-    if not token:
-        print("❌ No token received")
-        raise HTTPException(status_code=401, detail="No token provided")
+        # Remove 'Bearer ' prefix if present
+        if token.startswith("Bearer "):
+            token = token[7:]
 
-    try:
         payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
-        print("🧠 Token payload:", payload)
+        email: str = payload.get("sub")
+        if email is None:
+            logger.error("Token missing 'sub' claim")
+            raise credentials_exception
 
-        email = payload.get("sub")
-        if not email:
-            raise HTTPException(status_code=401, detail="Invalid token: missing subject")
     except JWTError as e:
-        print("❌ JWT decode error:", str(e))
-        raise HTTPException(status_code=401, detail="Could not validate token")
+        logger.error(f"JWT validation error: {str(e)}")
+        raise credentials_exception
 
     user = await users_collection.find_one({"email": email})
-    if not user:
-        raise HTTPException(status_code=404, detail="User not found")
+    if user is None:
+        logger.error(f"User not found for email: {email}")
+        raise credentials_exception
 
-    print("✅ Authenticated user:", user["email"])
-    return user
+    return user