labs/oauth/backend/app/services/token_service.py

"""Token management service"""

from typing import Optional, Dict, Any
from datetime import datetime, timedelta, timezone
from bson import ObjectId
from app.utils.database import get_database
from app.utils.security import create_access_token, create_refresh_token, decode_token
from app.models.user import Token
import logging
import secrets

logger = logging.getLogger(__name__)


class TokenService:
    """Service for handling token operations"""
    
    @staticmethod
    async def create_tokens(user_id: str, user_data: Dict[str, Any]) -> Token:
        """Create access and refresh tokens for a user"""
        # Create access token
        access_token = create_access_token(
            data={
                "sub": user_id,
                "username": user_data.get("username"),
                "role": user_data.get("role"),
                "email": user_data.get("email")
            }
        )
        
        # Create refresh token
        refresh_token = create_refresh_token(
            data={
                "sub": user_id,
                "username": user_data.get("username")
            }
        )
        
        # Store refresh token in database
        await TokenService.store_refresh_token(user_id, refresh_token)
        
        return Token(
            access_token=access_token,
            refresh_token=refresh_token,
            token_type="bearer"
        )
    
    @staticmethod
    async def store_refresh_token(user_id: str, token: str):
        """Store refresh token in database"""
        db = get_database()
        
        from app.config import settings
        expires_at = datetime.utcnow() + timedelta(
            days=settings.jwt_refresh_token_expire_days
        )
        
        token_doc = {
            "token": token,
            "user_id": user_id,
            "created_at": datetime.utcnow(),
            "expires_at": expires_at,
            "is_active": True
        }
        
        await db.refresh_tokens.insert_one(token_doc)
    
    @staticmethod
    async def verify_refresh_token(token: str) -> Optional[str]:
        """Verify refresh token and return user_id"""
        db = get_database()
        
        # Decode token
        payload = decode_token(token)
        if not payload or payload.get("type") != "refresh":
            return None
        
        # Check if token exists in database
        token_doc = await db.refresh_tokens.find_one({
            "token": token,
            "is_active": True
        })
        
        if not token_doc:
            return None
        
        # Check if token is expired
        if token_doc["expires_at"] < datetime.utcnow():
            # Mark token as inactive
            await db.refresh_tokens.update_one(
                {"_id": token_doc["_id"]},
                {"$set": {"is_active": False}}
            )
            return None
        
        return payload.get("sub")
    
    @staticmethod
    async def revoke_refresh_token(token: str):
        """Revoke a refresh token"""
        db = get_database()
        
        await db.refresh_tokens.update_one(
            {"token": token},
            {"$set": {"is_active": False}}
        )
    
    @staticmethod
    async def revoke_all_user_tokens(user_id: str):
        """Revoke all refresh tokens for a user"""
        db = get_database()
        
        await db.refresh_tokens.update_many(
            {"user_id": user_id},
            {"$set": {"is_active": False}}
        )
    
    @staticmethod
    async def create_authorization_code(
        user_id: str,
        client_id: str,
        redirect_uri: str,
        scope: str,
        state: Optional[str] = None
    ) -> str:
        """Create and store authorization code"""
        db = get_database()
        
        code = secrets.token_urlsafe(32)
        expires_at = datetime.utcnow() + timedelta(minutes=10)  # Code valid for 10 minutes
        
        code_doc = {
            "code": code,
            "user_id": user_id,
            "client_id": client_id,
            "redirect_uri": redirect_uri,
            "scope": scope,
            "state": state,
            "created_at": datetime.utcnow(),
            "expires_at": expires_at,
            "used": False
        }
        
        await db.authorization_codes.insert_one(code_doc)
        
        return code
    
    @staticmethod
    async def exchange_authorization_code(
        code: str,
        client_id: str,
        client_secret: str,
        redirect_uri: str
    ) -> Optional[Token]:
        """Exchange authorization code for tokens"""
        db = get_database()
        
        # Find and validate code
        code_doc = await db.authorization_codes.find_one({
            "code": code,
            "client_id": client_id,
            "redirect_uri": redirect_uri,
            "used": False
        })
        
        if not code_doc:
            return None
        
        # Check if code is expired
        if code_doc["expires_at"] < datetime.utcnow():
            return None
        
        # Validate client secret
        app = await db.applications.find_one({
            "client_id": client_id,
            "client_secret": client_secret
        })
        
        if not app:
            return None
        
        # Mark code as used
        await db.authorization_codes.update_one(
            {"_id": code_doc["_id"]},
            {"$set": {"used": True}}
        )
        
        # Get user
        user = await db.users.find_one({"_id": ObjectId(code_doc["user_id"])})
        if not user:
            return None
        
        # Create tokens
        return await TokenService.create_tokens(
            str(user["_id"]),
            {
                "username": user["username"],
                "role": user["role"],
                "email": user["email"]
            }
        )