solventum-image-metadata/backend/app/api/auth.py

"""
Authentication API Endpoints
Handles login, logout, token refresh, and Microsoft SSO.
"""

from fastapi import APIRouter, Depends, HTTPException, status, Request
from fastapi.responses import JSONResponse
from sqlalchemy.ext.asyncio import AsyncSession
from pydantic import BaseModel
from typing import Optional
import msal
import os

from app.core.database import get_db, UserRepository, AuditLogRepository
from app.core.auth import (
    verify_password,
    hash_password,
    create_tokens_response,
    verify_refresh_token,
    get_current_user_id
)
from app.core.redis_client import RedisSessionStore


router = APIRouter()


# ===== Request/Response Models =====

class LoginRequest(BaseModel):
    username: str
    password: str


class LoginResponse(BaseModel):
    access_token: str
    refresh_token: str
    token_type: str
    expires_in: int
    user: dict


class TokenRefreshRequest(BaseModel):
    refresh_token: str


class LogoutRequest(BaseModel):
    session_id: Optional[str] = None


# ===== Local Authentication Endpoints =====

@router.post("/login", response_model=LoginResponse)
async def login(
    login_data: LoginRequest,
    request: Request,
    db: AsyncSession = Depends(get_db)
):
    """
    Local authentication - username/password login.

    Returns JWT tokens + user info.
    """
    # Get user from database
    user = await UserRepository.get_by_username(db, login_data.username)

    # Validate user exists and password correct
    if not user or not user.password_hash:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid username or password"
        )

    if not verify_password(login_data.password, user.password_hash):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid username or password"
        )

    # Check if user is active
    if not user.is_active:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="User account is disabled"
        )

    # Create JWT tokens
    tokens = create_tokens_response(user.id)

    # Create user session in Redis
    redis: RedisSessionStore = request.app.state.redis
    session_id = await redis.create_user_session(
        user_id=user.id,
        refresh_token=tokens["refresh_token"],
        ip_address=request.client.host,
        user_agent=request.headers.get("user-agent", "")
    )

    # Update last login
    await UserRepository.update_last_login(db, user.id)

    # Log action
    await AuditLogRepository.log_action(
        db,
        user_id=user.id,
        action="login",
        details=f"Login from {request.client.host}"
    )

    return LoginResponse(
        **tokens,
        user=user.to_dict()
    )


@router.post("/token/refresh")
async def refresh_access_token(
    refresh_data: TokenRefreshRequest,
    request: Request,
    db: AsyncSession = Depends(get_db)
):
    """
    Refresh access token using refresh token.
    """
    # Verify refresh token
    try:
        user_id = verify_refresh_token(refresh_data.refresh_token)
    except HTTPException as e:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid refresh token"
        )

    # Check if user still exists and is active
    user = await UserRepository.get_by_id(db, user_id)
    if not user or not user.is_active:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="User not found or inactive"
        )

    # Create new tokens
    tokens = create_tokens_response(user.id)

    # Update Redis session with new refresh token
    redis: RedisSessionStore = request.app.state.redis
    # Note: We keep the old session_id but update the refresh token
    # In production, you might want to rotate session_id as well

    return {
        **tokens,
        "user": user.to_dict()
    }


@router.post("/logout")
async def logout(
    logout_data: LogoutRequest,
    request: Request,
    user_id: int = Depends(get_current_user_id),
    db: AsyncSession = Depends(get_db)
):
    """
    Logout user - invalidate session in Redis.
    """
    # Delete user session from Redis
    redis: RedisSessionStore = request.app.state.redis

    if logout_data.session_id:
        await redis.delete_user_session(logout_data.session_id)

    # Log action
    await AuditLogRepository.log_action(
        db,
        user_id=user_id,
        action="logout",
        details=f"Logout from {request.client.host}"
    )

    return {"message": "Logged out successfully"}


# ===== Microsoft SSO Endpoints =====

# Microsoft OAuth configuration
AZURE_CLIENT_ID = os.getenv("AZURE_CLIENT_ID")
AZURE_CLIENT_SECRET = os.getenv("AZURE_CLIENT_SECRET")
AZURE_TENANT_ID = os.getenv("AZURE_TENANT_ID")
REDIRECT_URI = os.getenv("REDIRECT_URI", "http://localhost:8000/auth/microsoft/callback")


@router.get("/microsoft/login")
async def microsoft_sso_login():
    """
    Redirect to Microsoft SSO login page.
    Returns auth URL for frontend to redirect to.
    """
    if not AZURE_CLIENT_ID or not AZURE_TENANT_ID:
        raise HTTPException(
            status_code=status.HTTP_501_NOT_IMPLEMENTED,
            detail="Microsoft SSO not configured"
        )

    # Create MSAL confidential client app
    msal_app = msal.ConfidentialClientApplication(
        client_id=AZURE_CLIENT_ID,
        client_credential=AZURE_CLIENT_SECRET,
        authority=f"https://login.microsoftonline.com/{AZURE_TENANT_ID}"
    )

    # Get authorization URL
    auth_url = msal_app.get_authorization_request_url(
        scopes=["User.Read"],
        redirect_uri=REDIRECT_URI
    )

    return {"auth_url": auth_url}


@router.get("/microsoft/callback")
async def microsoft_sso_callback(
    code: str,
    request: Request,
    db: AsyncSession = Depends(get_db)
):
    """
    Handle Microsoft SSO callback.
    Exchange authorization code for tokens and create user session.
    """
    if not AZURE_CLIENT_ID or not AZURE_TENANT_ID:
        raise HTTPException(
            status_code=status.HTTP_501_NOT_IMPLEMENTED,
            detail="Microsoft SSO not configured"
        )

    # Create MSAL app
    msal_app = msal.ConfidentialClientApplication(
        client_id=AZURE_CLIENT_ID,
        client_credential=AZURE_CLIENT_SECRET,
        authority=f"https://login.microsoftonline.com/{AZURE_TENANT_ID}"
    )

    # Acquire token by authorization code
    result = msal_app.acquire_token_by_authorization_code(
        code=code,
        scopes=["User.Read"],
        redirect_uri=REDIRECT_URI
    )

    if "access_token" not in result:
        error = result.get("error_description", "Unknown error")
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=f"SSO authentication failed: {error}"
        )

    # Get user info from Microsoft Graph
    import requests
    graph_response = requests.get(
        "https://graph.microsoft.com/v1.0/me",
        headers={"Authorization": f"Bearer {result['access_token']}"}
    )

    if graph_response.status_code != 200:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Failed to get user info from Microsoft Graph"
        )

    user_info = graph_response.json()

    # Create or update user in database
    username = user_info.get("userPrincipalName") or user_info.get("mail")
    email = user_info.get("mail")
    full_name = user_info.get("displayName")

    user = await UserRepository.get_by_username(db, username)

    if not user:
        # Create new SSO user
        user = await UserRepository.create_user(
            db,
            username=username,
            password_hash=None,  # SSO users don't have passwords
            email=email,
            full_name=full_name,
            auth_method="sso"
        )
    else:
        # Update existing user
        user.last_login = None  # Will be updated by update_last_login
        await db.commit()

    # Create JWT tokens
    tokens = create_tokens_response(user.id)

    # Create user session in Redis
    redis: RedisSessionStore = request.app.state.redis
    session_id = await redis.create_user_session(
        user_id=user.id,
        refresh_token=tokens["refresh_token"],
        ip_address=request.client.host,
        user_agent=request.headers.get("user-agent", "")
    )

    # Update last login
    await UserRepository.update_last_login(db, user.id)

    # Log action
    await AuditLogRepository.log_action(
        db,
        user_id=user.id,
        action="sso_login",
        details=f"SSO login from {request.client.host}"
    )

    return LoginResponse(
        **tokens,
        user=user.to_dict()
    )


# ===== User Info Endpoint =====

@router.get("/me")
async def get_current_user(
    user_id: int = Depends(get_current_user_id),
    db: AsyncSession = Depends(get_db)
):
    """
    Get current user info from JWT token.
    """
    user = await UserRepository.get_by_id(db, user_id)

    if not user:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="User not found"
        )

    return user.to_dict()


# ===== Admin Endpoints (for testing) =====

@router.post("/register")
async def register_user(
    login_data: LoginRequest,
    db: AsyncSession = Depends(get_db)
):
    """
    Register new user (for testing/development).
    In production, disable this or add admin auth.
    """
    # Check if user already exists
    existing_user = await UserRepository.get_by_username(db, login_data.username)
    if existing_user:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Username already exists"
        )

    # Create new user
    password_hashed = hash_password(login_data.password)
    user = await UserRepository.create_user(
        db,
        username=login_data.username,
        password_hash=password_hashed,
        email=None,
        full_name=None,
        auth_method="local"
    )

    return {
        "message": "User created successfully",
        "user": user.to_dict()
    }