gateway/modules/routes/routeSecurityMsft.py

"""
Routes for Microsoft authentication.
"""

from fastapi import APIRouter, HTTPException, Request, Response, status, Depends, Body, Query
from fastapi.responses import HTMLResponse, RedirectResponse, JSONResponse
import logging
import json
from typing import Dict, Any, Optional
from datetime import datetime, timedelta
import msal
import httpx

from modules.shared.configuration import APP_CONFIG
from modules.interfaces.interfaceAppObjects import getInterface, getRootInterface
from modules.interfaces.interfaceAppModel import AuthAuthority, User, Token, ConnectionStatus, UserConnection
from modules.security.auth import getCurrentUser, limiter, createAccessToken
from modules.shared.attributeUtils import ModelMixin
from modules.shared.timezoneUtils import get_utc_now, create_expiration_timestamp, get_utc_timestamp

# Configure logger
logger = logging.getLogger(__name__)

# Create router
router = APIRouter(
    prefix="/api/msft",
    tags=["Security Microsoft"],
    responses={
        404: {"description": "Not found"}, 
        400: {"description": "Bad request"},
        401: {"description": "Unauthorized"},
        403: {"description": "Forbidden"},
        500: {"description": "Internal server error"}
    }
)

# Microsoft OAuth configuration
CLIENT_ID = APP_CONFIG.get("Service_MSFT_CLIENT_ID")
CLIENT_SECRET = APP_CONFIG.get("Service_MSFT_CLIENT_SECRET")
TENANT_ID = APP_CONFIG.get("Service_MSFT_TENANT_ID", "common")
REDIRECT_URI = APP_CONFIG.get("Service_MSFT_REDIRECT_URI")
AUTHORITY = f"https://login.microsoftonline.com/{TENANT_ID}"
SCOPES = [
    "Mail.ReadWrite",           # Read and write mail
    "Mail.Send",                # Send mail
    "Mail.ReadWrite.Shared",    # Access shared mailboxes
    "User.Read"                 # Read user profile
]

@router.get("/login")
@limiter.limit("5/minute")
async def login(
    request: Request,
    state: str = Query("login", description="State parameter to distinguish between login and connection flows"),
    connectionId: Optional[str] = Query(None, description="Connection ID for connection flow")
) -> RedirectResponse:
    """Initiate Microsoft login"""
    try:
        # Create MSAL app
        msal_app = msal.ConfidentialClientApplication(
            CLIENT_ID,
            authority=AUTHORITY,
            client_credential=CLIENT_SECRET
        )
        
        # Generate auth URL with state - use state as is if it's already JSON, otherwise create new state
        try:
            # Try to parse state as JSON to check if it's already encoded
            json.loads(state)
            state_param = state  # Use state as is if it's valid JSON
        except json.JSONDecodeError:
            # If not JSON, create new state object
            state_param = json.dumps({
                "type": state,
                "connectionId": connectionId
            })
        
        # MSAL automatically adds openid, profile, offline_access - we just need to provide our business scopes
        auth_url = msal_app.get_authorization_request_url(
            scopes=SCOPES,  # Only our business scopes - MSAL adds the required ones automatically
            redirect_uri=REDIRECT_URI,
            state=state_param,
            prompt="select_account"  # Force account selection screen
        )
        
        return RedirectResponse(auth_url)
        
    except Exception as e:
        logger.error(f"Error initiating Microsoft login: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to initiate Microsoft login: {str(e)}"
        )

@router.get("/auth/callback")
async def auth_callback(code: str, state: str, request: Request) -> HTMLResponse:
    """Handle Microsoft OAuth callback"""
    try:
        # Parse state
        state_data = json.loads(state)
        state_type = state_data.get("type", "login")
        connection_id = state_data.get("connectionId")
        user_id = state_data.get("userId")  # Get user ID from state
        
        logger.info(f"Processing Microsoft auth callback: state_type={state_type}, connection_id={connection_id}, user_id={user_id}")
        
        # Create MSAL app
        msal_app = msal.ConfidentialClientApplication(
            CLIENT_ID,
            authority=AUTHORITY,
            client_credential=CLIENT_SECRET
        )
        
        # Get token from code - MSAL automatically handles the required scopes
        token_response = msal_app.acquire_token_by_authorization_code(
            code,
            scopes=SCOPES,  # Only our business scopes - MSAL adds the required ones automatically
            redirect_uri=REDIRECT_URI
        )
        
        if "error" in token_response:
            logger.error(f"Token acquisition failed: {token_response['error']}")
            return HTMLResponse(
                content="<html><body><h1>Authentication Failed</h1><p>Could not acquire token.</p></body></html>",
                status_code=400
            )
        
        # Get user info using the access token
        headers = {
            'Authorization': f"Bearer {token_response['access_token']}",
            'Content-Type': 'application/json'
        }
        async with httpx.AsyncClient() as client:
            user_info_response = await client.get(
                "https://graph.microsoft.com/v1.0/me",
                headers=headers
            )
            if user_info_response.status_code != 200:
                logger.error(f"Failed to get user info: {user_info_response.text}")
                raise HTTPException(
                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                    detail="Failed to get user info from Microsoft"
                )
            user_info = user_info_response.json()
            logger.info(f"Got user info from Microsoft: {user_info.get('userPrincipalName')}")
        
        if state_type == "login":
            # Handle login flow
            rootInterface = getRootInterface()
            user = rootInterface.getUserByUsername(user_info.get("userPrincipalName"))
            
            if not user:
                logger.info(f"Creating new user for {user_info.get('userPrincipalName')}")
                # Create new user if doesn't exist
                user = rootInterface.createUser(
                    username=user_info.get("userPrincipalName"),
                    email=user_info.get("mail"),
                    fullName=user_info.get("displayName"),
                    authenticationAuthority=AuthAuthority.MSFT,
                    externalId=user_info.get("id"),
                    externalUsername=user_info.get("userPrincipalName"),
                    externalEmail=user_info.get("mail")
                )
            
            # Create token
            token = Token(
                userId=user.id,  # Use local user's ID
                authority=AuthAuthority.MSFT,
                tokenAccess=token_response["access_token"],
                tokenRefresh=token_response.get("refresh_token", ""),
                tokenType=token_response.get("token_type", "bearer"),
                expiresAt=create_expiration_timestamp(token_response.get("expires_in", 0)),
                createdAt=get_utc_timestamp()
            )
            
            # Save access token (no connectionId)
            appInterface = getInterface(user)
            appInterface.saveAccessToken(token)

            # Create JWT token data
            jwt_token_data = {
                "sub": user.username,
                "mandateId": str(user.mandateId),
                "userId": str(user.id),
                "authenticationAuthority": AuthAuthority.MSFT
            }

            # Create JWT access token
            jwt_token, jwt_expires_at = createAccessToken(jwt_token_data)

            # Create JWT token
            jwt_token_obj = Token(
                userId=user.id,
                authority=AuthAuthority.MSFT,
                tokenAccess=jwt_token,
                tokenType="bearer",
                expiresAt=jwt_expires_at.timestamp(),
                createdAt=get_utc_timestamp()
            )
            
            # Save JWT access token
            appInterface.saveAccessToken(jwt_token_obj)
            
            # Convert token to dict and ensure proper timestamp handling
            token_dict = jwt_token_obj.to_dict()
            # Remove datetime conversion logic - models now handle this automatically
            # The token model already returns float timestamps
            
            # Return success page with token data
            return HTMLResponse(
                content=f"""
                <html>
                    <head><title>Authentication Successful</title></head>
                    <body>
                        <script>
                            if (window.opener) {{
                                window.opener.postMessage({{
                                    type: 'msft_auth_success',
                                    token_data: {json.dumps(token_dict)}
                                }}, '*');
                            }}
                            setTimeout(() => window.close(), 1000);
                        </script>
                    </body>
                </html>
                """
            )
        else:
            # Handle connection flow
            if not connection_id or not user_id:
                logger.error("Connection ID or User ID is missing in connection flow")
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="Connection ID and User ID are required for connection flow"
                )
            
            # Get user directly by ID
            rootInterface = getRootInterface()
            user = rootInterface.getUser(user_id)
            
            if not user:
                logger.error(f"User {user_id} not found in database")
                return HTMLResponse(
                    content=f"""
                    <html>
                        <head><title>Connection Failed</title></head>
                        <body>
                            <script>
                                if (window.opener) {{
                                    window.opener.postMessage({{
                                        type: 'msft_connection_error',
                                        error: 'User not found in database'
                                    }}, '*');
                                    // Wait for message to be sent before closing
                                    setTimeout(() => window.close(), 1000);
                                }} else {{
                                    window.close();
                                }}
                            </script>
                        </body>
                    </html>
                    """,
                    status_code=404
                )
            
            # Get the connection from the connections table
            interface = getInterface(user)
            connections = interface.getUserConnections(user_id)
            connection = None
            for conn in connections:
                if conn.id == connection_id:
                    connection = conn
                    logger.info(f"Found existing connection for user {user.username}")
                    break
            
            try:
                if not connection:
                    logger.error(f"Connection {connection_id} not found in user's connections")
                    return HTMLResponse(
                        content=f"""
                        <html>
                            <head><title>Connection Failed</title></head>
                            <body>
                                <script>
                                    if (window.opener) {{
                                        window.opener.postMessage({{
                                            type: 'msft_connection_error',
                                            error: 'Connection not found in user\'s connections'
                                        }}, '*');
                                        // Wait for message to be sent before closing
                                        setTimeout(() => window.close(), 1000);
                                    }} else {{
                                        window.close();
                                    }}
                                </script>
                            </body>
                        </html>
                        """,
                        status_code=404
                    )
                
                logger.info(f"Updating connection {connection_id} for user {user.username}")
                # Update connection with external service details
                connection.status = ConnectionStatus.ACTIVE
                connection.lastChecked = get_utc_timestamp()
                connection.expiresAt = get_utc_timestamp() + token_response.get("expires_in", 0)
                connection.externalId = user_info.get("id")
                connection.externalUsername = user_info.get("userPrincipalName")
                connection.externalEmail = user_info.get("mail")
                
                # Update connection record directly
                rootInterface.db.recordModify("connections", connection_id, connection.to_dict())
                
                # Clear cache to ensure fresh data
                rootInterface.db.clearTableCache("connections")
                
                # Save token
        
                token = Token(
                    userId=user.id,  # Use local user's ID
                    authority=AuthAuthority.MSFT,
                    connectionId=connection_id,  # Link token to this specific connection
                    tokenAccess=token_response["access_token"],
                    tokenRefresh=token_response.get("refresh_token", ""),
                    tokenType=token_response.get("token_type", "bearer"),
                    expiresAt=create_expiration_timestamp(token_response.get("expires_in", 0)),
                    createdAt=get_utc_timestamp()
                )
                
        
                interface.saveConnectionToken(token)
        
                
                # Return success page with connection data
                return HTMLResponse(
                    content=f"""
                    <html>
                        <head><title>Connection Successful</title></head>
                        <body>
                            <script>
                                if (window.opener) {{
                                    window.opener.postMessage({{
                                        type: 'msft_connection_success',
                                        connection: {{
                                            id: '{connection.id}',
                                            status: 'connected',
                                            type: 'msft',
                                            lastChecked: {get_utc_timestamp()},
                                            expiresAt: {create_expiration_timestamp(token_response.get("expires_in", 0))}
                                        }}
                                    }}, '*');
                                    // Wait for message to be sent before closing
                                    setTimeout(() => window.close(), 1000);
                                }} else {{
                                    window.close();
                                }}
                            </script>
                        </body>
                    </html>
                    """
                )
            except Exception as e:
                logger.error(f"Error updating connection or saving token: {str(e)}", exc_info=True)
                return HTMLResponse(
                    content=f"""
                    <html>
                        <head><title>Connection Failed</title></head>
                        <body>
                            <script>
                                if (window.opener) {{
                                    window.opener.postMessage({{
                                        type: 'msft_connection_error',
                                        error: 'Failed to update connection or save token: {str(e)}'
                                    }}, '*');
                                    // Wait for message to be sent before closing
                                    setTimeout(() => window.close(), 1000);
                                }} else {{
                                    window.close();
                                }}
                            </script>
                        </body>
                    </html>
                    """,
                    status_code=500
                )
            
    except Exception as e:
        logger.error(f"Error in auth callback: {str(e)}", exc_info=True)
        return HTMLResponse(
            content=f"""
            <html>
                <head><title>Authentication Failed</title></head>
                <body>
                    <script>
                        if (window.opener) {{
                            window.opener.postMessage({{
                                type: 'msft_connection_error',
                                error: 'Authentication failed: {str(e)}'
                            }}, '*');
                            // Wait for message to be sent before closing
                            setTimeout(() => window.close(), 1000);
                        }} else {{
                            window.close();
                        }}
                    </script>
                </body>
            </html>
            """,
            status_code=500
        )

@router.get("/me", response_model=User)
@limiter.limit("30/minute")
async def get_current_user(
    request: Request,
    currentUser: User = Depends(getCurrentUser)
) -> User:
    """Get current user information"""
    try:
        return currentUser
    except Exception as e:
        logger.error(f"Error getting current user: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to get current user: {str(e)}"
        )

@router.post("/logout")
@limiter.limit("10/minute")
async def logout(
    request: Request,
    currentUser: User = Depends(getCurrentUser)
) -> Dict[str, Any]:
    """Logout current user"""
    try:
        appInterface = getInterface(currentUser)
        appInterface.logout()
        return {"message": "Logged out successfully"}
    except Exception as e:
        logger.error(f"Error during logout: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to logout: {str(e)}"
        )

@router.post("/cleanup")
@limiter.limit("5/minute")
async def cleanup_expired_tokens(
    request: Request,
    currentUser: User = Depends(getCurrentUser)
) -> Dict[str, Any]:
    """Clean up expired tokens for the current user"""
    try:
        appInterface = getInterface(currentUser)
        
        # Clean up expired tokens
        cleaned_count = appInterface.cleanupExpiredTokens()
        
        return {
            "message": f"Cleanup completed successfully",
            "tokens_cleaned": cleaned_count
        }
        
    except Exception as e:
        logger.error(f"Error cleaning up expired tokens: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to cleanup expired tokens: {str(e)}"
        )

@router.post("/refresh")
@limiter.limit("10/minute")
async def refresh_token(
    request: Request,
    currentUser: User = Depends(getCurrentUser)
) -> Dict[str, Any]:
    """Refresh Microsoft OAuth token for current user"""
    try:
        appInterface = getInterface(currentUser)
        
        # Find Microsoft connection for this user
        logger.debug(f"Looking for Microsoft connection for user {currentUser.id}")
        
        connections = appInterface.getUserConnections(currentUser.id)
        msft_connection = None
        
        for conn in connections:
            if conn.authority == AuthAuthority.MSFT:
                msft_connection = conn
                break
        
        if not msft_connection:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="No Microsoft connection found for current user"
            )
        
        logger.debug(f"Found Microsoft connection: {msft_connection.id}, status={msft_connection.status}")
        
        # Get the token for this specific connection using the new method
        # Enable auto-refresh to handle expired tokens gracefully
        current_token = appInterface.getConnectionToken(msft_connection.id, auto_refresh=True)
        
        if not current_token:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="No Microsoft token found for this connection"
            )
        

        
        # Always attempt refresh (as per your requirement)
        from modules.security.tokenManager import TokenManager
        token_manager = TokenManager()
        
        refreshed_token = token_manager.refresh_token(current_token)
        if refreshed_token:
            # Save the new connection token (which will automatically replace old ones)
            appInterface.saveConnectionToken(refreshed_token)
            
            # Update the connection's expiration time
            msft_connection.expiresAt = float(refreshed_token.expiresAt)
            msft_connection.lastChecked = get_utc_timestamp()
            msft_connection.status = ConnectionStatus.ACTIVE
            
            # Save updated connection
            appInterface.db.recordModify("connections", msft_connection.id, msft_connection.to_dict())
            
            # Calculate time until expiration
            current_time = get_utc_timestamp()
            expires_in = int(refreshed_token.expiresAt - current_time)
            
            return {
                "message": "Token refreshed successfully",
                "expires_at": refreshed_token.expiresAt,
                "expires_in_seconds": expires_in
            }
        else:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail="Failed to refresh token"
            )
            
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error refreshing Microsoft token: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to refresh token: {str(e)}"
        )