import bcrypt import os from datetime import datetime, timedelta from functools import wraps from typing import Any, Dict, Union from jose import jwt, JWTError from fastapi import Depends, HTTPException, Request, status from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials from sqlalchemy import select, desc from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.exc import SQLAlchemyError, DBAPIError from ..database.db import session_pool from ..errors import DatabaseConnectionError from ..model.user import User from ..model.profile import Profile from ..model.login_session import LoginSession from ..utils.token import decodeJWT from ..core.logger import SingletonLogger def hash_password(password: str) -> str: """Hash a password using bcrypt""" return bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt()).decode("utf-8") def verify_password(plain_password: str, hashed_password: str) -> bool: """Verify a password against its hash""" return bcrypt.checkpw( plain_password.encode("utf-8"), hashed_password.encode("utf-8") ) async def build_token_payload(user: User, session: AsyncSession) -> Dict[str, Any]: """ Build JWT token payload with user and profile information. Args: user: User object session: Database session to fetch profile Returns: Dict containing user data (excluding password) """ # Get user profile profile_result = await session.execute( select(Profile).where(Profile.user_id == user.id) ) profile = profile_result.scalar_one_or_none() payload = { "sub": str(user.id), "user_id": user.id, "username": user.username, "full_name": user.full_name, "email": user.email, "google_id": user.google_id, "profile": ( { "phone": profile.phone if profile else None, "avatar": profile.avatar if profile else None, "bio": profile.bio if profile else None, } if profile else None ), } return payload def create_access_token( data: dict, expires_delta: timedelta | None = None ) -> tuple[str, datetime]: """Create JWT access token with user and profile information""" to_encode = data.copy() expire = datetime.utcnow() + ( expires_delta or timedelta(minutes=int(os.getenv("JWT_ACCESS_TOKEN_EXPIRE_MINUTES", 30))) ) to_encode.update({"exp": expire}) encoded_jwt = jwt.encode( to_encode, os.getenv("JWT_SECRET_KEY"), algorithm=os.getenv("JWT_ALGORITHM") ) return encoded_jwt, expire def create_refresh_token( subject: Union[str, Any], expires_delta: timedelta | None = None ) -> tuple[str, datetime]: """Create JWT refresh token""" if expires_delta is not None: expires_delta_time = datetime.utcnow() + expires_delta else: expires_delta_time = datetime.utcnow() + timedelta( minutes=int(os.getenv("JWT_REFRESH_TOKEN_EXPIRE_MINUTES", 10080)) ) to_encode = {"exp": expires_delta_time, "sub": str(subject)} encoded_jwt = jwt.encode( to_encode, os.getenv("JWT_SECRET_KEY"), os.getenv("JWT_ALGORITHM") ) return encoded_jwt, expires_delta_time class JWTBearer(HTTPBearer): """Custom JWT Bearer authentication class with token verification.""" def __init__(self, auto_error: bool = True): super(JWTBearer, self).__init__(auto_error=auto_error) async def __call__(self, request: Request): credentials: HTTPAuthorizationCredentials = await super( JWTBearer, self ).__call__(request) if credentials: logger = getattr( request.app.state, "logger", SingletonLogger().get_logger() ) if not credentials.scheme == "Bearer": logger.error("Invalid authentication scheme.") raise HTTPException( status_code=403, detail="Invalid authentication scheme." ) if not self.verify_jwt(credentials.credentials): logger.error("Invalid token or expired token.") raise HTTPException( status_code=403, detail="Invalid token or expired token." ) return credentials.credentials else: logger = getattr( request.app.state, "logger", SingletonLogger().get_logger() ) logger.error("Invalid authorization code.") raise HTTPException(status_code=403, detail="Invalid authorization code.") def verify_jwt(self, jwtoken: str) -> bool: """Verify if the JWT token is valid.""" isTokenValid: bool = False try: payload = decodeJWT(jwtoken) except Exception as e: SingletonLogger().get_logger().error(f"JWT verification error: {e}") payload = None if payload: isTokenValid = True return isTokenValid async def get_current_user( credentials: HTTPAuthorizationCredentials = Depends(HTTPBearer()), ) -> int: """ Dependency function to get current user ID from JWT token. Validates token against the login_session table to ensure it's still active. Returns the user ID for use in route dependencies. """ credentials_exception = HTTPException( status_code=401, detail="Could not validate credentials", headers={"WWW-Authenticate": "Bearer"}, ) try: payload = jwt.decode( credentials.credentials, os.getenv("JWT_SECRET_KEY"), algorithms=[os.getenv("JWT_ALGORITHM")], ) user_id: str = payload.get("sub") if user_id is None: raise credentials_exception except JWTError: SingletonLogger().get_logger().error("JWT decode error") raise credentials_exception # Verify token is active in login_session table try: async with session_pool() as session: session_result = await session.execute( select(LoginSession) .filter_by( user_id=int(user_id), access_token=credentials.credentials, is_active=True, ) .order_by(desc(LoginSession.created_at)) ) session_record = session_result.scalar_one_or_none() if not session_record: SingletonLogger().get_logger().error( f"Token not found in login_session or inactive for user {user_id}" ) raise HTTPException( status_code=401, detail="Token is not active or has been revoked", headers={"WWW-Authenticate": "Bearer"}, ) except (DBAPIError, SQLAlchemyError) as db_err: SingletonLogger().get_logger().exception( f"Database connection error while validating token for user {user_id}: {db_err}" ) raise DatabaseConnectionError(str(db_err)) return int(user_id) def token_required(func): """ Decorator that verifies the JWT token and checks if user is logged in. Validates token against the login_session table to ensure it's still active. Injects user_id into the function kwargs. Usage: @router.get("/protected") @token_required async def protected_route(user_id: int): # user_id is automatically injected return {"user_id": user_id} """ @wraps(func) async def wrapper(*args, **kwargs): # Get request from args (first argument in FastAPI route handlers) request = None for arg in args: if isinstance(arg, Request): request = arg break # Prepare logger (from request app.state if available) logger = ( getattr(request.app.state, "logger", SingletonLogger().get_logger()) if request else SingletonLogger().get_logger() ) # Extract Authorization header auth_header = None if request: auth_header = request.headers.get("Authorization") if not auth_header or not auth_header.startswith("Bearer "): logger.error("No authorization header or invalid format.") raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Not authenticated", headers={"WWW-Authenticate": "Bearer"}, ) token = auth_header.replace("Bearer ", "") # Decode the token payload = decodeJWT(token) if not payload: logger.error("Invalid token.") raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token.", headers={"WWW-Authenticate": "Bearer"}, ) user_id = payload.get("sub") if not user_id: logger.error("Invalid token payload.") raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token payload.", headers={"WWW-Authenticate": "Bearer"}, ) # Verify token exists in login_session and is active try: async with session_pool() as session: result = await session.execute( select(LoginSession) .filter_by(user_id=int(user_id), access_token=token, is_active=True) .order_by(desc(LoginSession.created_at)) ) session_record = result.scalar_one_or_none() if not session_record: logger.error("Invalid or inactive access token.") raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid or inactive access token.", headers={"WWW-Authenticate": "Bearer"}, ) except (DBAPIError, SQLAlchemyError) as db_err: logger.exception( f"Database connection error while checking token for user {user_id}: {db_err}" ) raise DatabaseConnectionError(str(db_err)) # Inject user_id into kwargs kwargs["user_id"] = int(user_id) return await func(*args, **kwargs) return wrapper async def store_token_in_session( session: AsyncSession, user_id: int, access_token: str, refresh_token: str | None = None, login_method: str = "password", device_info: str | None = None, ip_address: str | None = None, user_agent: str | None = None, token_expires_at: datetime | None = None, ) -> LoginSession: """ Store JWT tokens in the login_session table. Args: session: Database session user_id: ID of the user access_token: JWT access token refresh_token: JWT refresh token (optional) login_method: Method of login ('password', 'google', etc.) device_info: Device information (optional) ip_address: IP address (optional) user_agent: User agent string (optional) token_expires_at: Token expiration time (optional) Returns: LoginSession: The created login session record """ login_session = LoginSession( user_id=user_id, login_method=login_method, is_active=True, access_token=access_token, refresh_token=refresh_token, device_info=device_info, ip_address=ip_address, user_agent=user_agent, token_expires_at=token_expires_at, ) session.add(login_session) await session.commit() await session.refresh(login_session) SingletonLogger().get_logger().info( f"Token stored in login_session for user {user_id}" ) return login_session async def revoke_token(session: AsyncSession, access_token: str) -> bool: """ Revoke a JWT token by setting the login session to inactive. Args: session: Database session access_token: The access token to revoke Returns: bool: True if token was revoked, False if not found """ result = await session.execute( select(LoginSession).filter_by(access_token=access_token, is_active=True) ) session_record = result.scalar_one_or_none() if session_record: session_record.is_active = False session_record.logout_at = datetime.utcnow() await session.commit() SingletonLogger().get_logger().info( f"Token revoked for user {session_record.user_id}" ) return True SingletonLogger().get_logger().warning("Token not found or already revoked") return False async def revoke_all_user_tokens(session: AsyncSession, user_id: int) -> int: """ Revoke all active tokens for a specific user. Args: session: Database session user_id: ID of the user Returns: int: Number of tokens revoked """ result = await session.execute( select(LoginSession).filter_by(user_id=user_id, is_active=True) ) sessions = result.scalars().all() count = 0 for login_session in sessions: login_session.is_active = False login_session.logout_at = datetime.utcnow() count += 1 await session.commit() SingletonLogger().get_logger().info(f"Revoked {count} sessions for user {user_id}") return count