gateway/routes/routeMsft.py

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

from modules.auth import getCurrentActiveUser, getUserContext, createAccessToken, ACCESS_TOKEN_EXPIRE_MINUTES
from modules.configuration import APP_CONFIG
from modules.lucydomInterface import getLucydomInterface
from modules.gatewayInterface import getGatewayInterface

# Configure logger
logger = logging.getLogger(__name__)

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

# Azure AD configuration - load from config
CLIENT_ID = APP_CONFIG.get("Agent_Mail_MSFT_CLIENT_ID")
CLIENT_SECRET = APP_CONFIG.get("Agent_Mail_MSFT_CLIENT_SECRET")
TENANT_ID = APP_CONFIG.get("Agent_Mail_MSFT_TENANT_ID", "common")  # Use 'common' for multi-tenant
AUTHORITY = f"https://login.microsoftonline.com/{TENANT_ID}"
SCOPES = ["Mail.ReadWrite", "User.Read"]
REDIRECT_URI = APP_CONFIG.get("Agent_Mail_MSFT_REDIRECT_URI")

# Initialize MSAL application
app_config = {
    "client_id": CLIENT_ID,
    "client_credential": CLIENT_SECRET,
    "authority": AUTHORITY,
    "redirect_uri": REDIRECT_URI
}

async def save_token_to_file(token_data, currentUser: Dict[str, Any]):
    """Save token data to database using LucyDOMInterface"""
    try:
        # Get current user context
        mandateId, userId = await getUserContext(currentUser)
        if not mandateId or not userId:
            logger.error("No user context available for token storage")
            return False
            
        # Get LucyDOM interface for current user
        mydom = getLucydomInterface(
            mandateId=mandateId,
            userId=userId
        )
        if not mydom:
            logger.error("No LucyDOM interface available for token storage")
            return False
            
        # Save token to database
        success = mydom.saveMsftToken(token_data)
        if success:
            logger.info("Token saved successfully to database")
            return True
        else:
            logger.error("Failed to save token to database")
            return False
            
    except Exception as e:
        logger.error(f"Error saving token: {str(e)}")
        return False

async def load_token_from_file(currentUser: Dict[str, Any]):
    """Load token data from database using LucyDOMInterface"""
    try:
        # Get current user context
        mandateId, userId = await getUserContext(currentUser)
        if not mandateId or not userId:
            logger.error("No user context available for token retrieval")
            return None
            
        # Get LucyDOM interface for current user
        mydom = getLucydomInterface(
            mandateId=mandateId,
            userId=userId
        )
        if not mydom:
            logger.error("No LucyDOM interface available for token retrieval")
            return None
            
        # Get token from database
        token_data = mydom.getMsftToken()
        if token_data:
            logger.info("Token loaded successfully from database")
            return token_data
        else:
            logger.info("No token found in database")
            return None
            
    except Exception as e:
        logger.error(f"Error loading token: {str(e)}")
        return None

def get_user_info_from_token(access_token: str) -> Optional[Dict[str, Any]]:
    """Get user information using the access token"""
    import requests
    headers = {
        'Authorization': f'Bearer {access_token}',
        'Content-Type': 'application/json'
    }
    
    try:
        response = requests.get('https://graph.microsoft.com/v1.0/me', headers=headers)
        if response.status_code == 200:
            user_data = response.json()
            return {
                "name": user_data.get("displayName", ""),
                "email": user_data.get("userPrincipalName", ""),
                "id": user_data.get("id", "")
            }
        else:
            logger.error(f"Error getting user info: {response.status_code} - {response.text}")
            return None
    except Exception as e:
        logger.error(f"Exception getting user info: {str(e)}")
        return None

def verify_token(token: str) -> bool:
    """Verify the access token is valid"""
    import requests
    headers = {
        'Authorization': f'Bearer {token}',
        'Content-Type': 'application/json'
    }
    
    try:
        logger.info("Verifying token validity...")
        response = requests.get('https://graph.microsoft.com/v1.0/me', headers=headers)
        
        if response.status_code == 200:
            logger.info("Token verification successful")
            return True
        else:
            logger.error(f"Token verification failed: {response.status_code} - {response.text}")
            return False
    except Exception as e:
        logger.error(f"Exception verifying token: {str(e)}")
        return False

async def refresh_token(user_id: str, currentUser: Dict[str, Any]) -> bool:
    """Refresh the access token using the stored refresh token"""
    token_data = await load_token_from_file(currentUser)
    if not token_data or not token_data.get("refresh_token"):
        logger.warning("No refresh token available")
        return False
    
    msal_app = msal.ConfidentialClientApplication(
        app_config["client_id"],
        authority=app_config["authority"],
        client_credential=app_config["client_credential"]
    )
    
    result = msal_app.acquire_token_by_refresh_token(
        token_data["refresh_token"],
        scopes=SCOPES
    )
    
    if "error" in result:
        logger.error(f"Error refreshing token: {result.get('error')}")
        return False
    
    # Update tokens in storage
    token_data["access_token"] = result["access_token"]
    if "refresh_token" in result:
        token_data["refresh_token"] = result["refresh_token"]
    
    await save_token_to_file(token_data, currentUser)
    logger.info("Access token refreshed successfully")
    return True

@router.get("/login")
async def login():
    """Initiate Microsoft login for the current user"""
    try:
        # Create a confidential client application
        msal_app = msal.ConfidentialClientApplication(
            app_config["client_id"],
            authority=app_config["authority"],
            client_credential=app_config["client_credential"]
        )
        
        # Build the auth URL with a random state
        state = secrets.token_urlsafe(32)
        
        auth_url = msal_app.get_authorization_request_url(
            SCOPES,
            state=state,  # Use random state
            redirect_uri=app_config["redirect_uri"]
        )
        
        logger.info(f"Redirecting to Microsoft login")
        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):
    """Handle Microsoft OAuth callback"""
    try:
        # Create MSAL app instance
        app = msal.ConfidentialClientApplication(
            client_id=CLIENT_ID,
            client_credential=CLIENT_SECRET,
            authority=AUTHORITY
        )
        
        # Exchange code for token
        token_response = app.acquire_token_by_authorization_code(
            code=code,
            scopes=SCOPES,
            redirect_uri=REDIRECT_URI
        )
        
        if "error" in token_response:
            logger.error(f"Token acquisition failed: {token_response['error']}")
            return HTMLResponse(
                content="""
                <html>
                    <head>
                        <title>Authentication Failed</title>
                        <style>
                            body { font-family: Arial, sans-serif; text-align: center; margin-top: 50px; }
                            .error { color: red; }
                        </style>
                    </head>
                    <body>
                        <h1 class="error">Authentication Failed</h1>
                        <p>Please try again.</p>
                        <script>
                            setTimeout(() => window.close(), 3000);
                        </script>
                    </body>
                </html>
                """,
                status_code=400
            )
        
        # Get user info from token
        user_info = get_user_info_from_token(token_response["access_token"])
        if not user_info:
            logger.error("Failed to get user info from token")
            return HTMLResponse(
                content="""
                <html>
                    <head>
                        <title>Authentication Failed</title>
                        <style>
                            body { font-family: Arial, sans-serif; text-align: center; margin-top: 50px; }
                            .error { color: red; }
                        </style>
                    </head>
                    <body>
                        <h1 class="error">Authentication Failed</h1>
                        <p>Could not retrieve user information.</p>
                        <script>
                            setTimeout(() => window.close(), 3000);
                        </script>
                    </body>
                </html>
                """,
                status_code=400
            )
        
        # Add user info to token data
        token_response["user_info"] = user_info
        
        # Store tokens in session storage for the frontend to pick up
        response = HTMLResponse(
            content=f"""
            <html>
                <head>
                    <title>Authentication Successful</title>
                    <style>
                        body {{ font-family: Arial, sans-serif; text-align: center; margin-top: 50px; }}
                        .success {{ color: green; }}
                    </style>
                </head>
                <body>
                    <h1 class="success">Authentication Successful</h1>
                    <p>Welcome, {user_info.get('name', 'User')}!</p>
                    <p>This window will close automatically.</p>
                    <script>
                        // Store token data in session storage
                        sessionStorage.setItem('msft_token_data', JSON.stringify({json.dumps(token_response)}));
                        
                        // Notify parent window of success
                        if (window.opener) {{
                            window.opener.postMessage({{ 
                                type: 'msft_auth_success', 
                                user: {json.dumps(user_info)},
                                token_data: {json.dumps(token_response)}
                            }}, '*');
                        }}
                        // Close window after 3 seconds
                        setTimeout(() => window.close(), 3000);
                    </script>
                </body>
            </html>
            """
        )
        
        return response
        
    except Exception as e:
        logger.error(f"Authentication failed: {str(e)}")
        return HTMLResponse(
            content="""
            <html>
                <head>
                    <title>Authentication Failed</title>
                    <style>
                        body { font-family: Arial, sans-serif; text-align: center; margin-top: 50px; }
                        .error { color: red; }
                    </style>
                </head>
                <body>
                    <h1 class="error">Authentication Failed</h1>
                    <p>An error occurred during authentication.</p>
                    <script>
                        setTimeout(() => window.close(), 3000);
                    </script>
                </body>
            </html>
            """,
            status_code=500
        )

@router.get("/status")
async def auth_status(currentUser: Dict[str, Any] = Depends(getCurrentActiveUser)):
    """Check Microsoft authentication status"""
    try:
        # Get current user context
        mandateId, userId = await getUserContext(currentUser)
        if not mandateId or not userId:
            logger.info("No user context found")
            return JSONResponse({
                "authenticated": False,
                "message": "Not authenticated with Microsoft"
            })
        
        # Check if we have a token for the current user
        token_data = await load_token_from_file(currentUser)
        
        if not token_data:
            logger.info(f"No token data found for user {userId}")
            return JSONResponse({
                "authenticated": False,
                "message": "Not authenticated with Microsoft"
            })
        
        # Verify token is still valid
        if not verify_token(token_data["access_token"]):
            logger.info("Token invalid, attempting refresh")
            # Try to refresh the token
            if not await refresh_token(userId, currentUser):
                logger.info("Token refresh failed")
                return JSONResponse({
                    "authenticated": False,
                    "message": "Token expired and refresh failed"
                })
            # Reload token data after refresh
            token_data = await load_token_from_file(currentUser)
        
        # Get user info from token data
        user_info = token_data.get("user_info")
        if not user_info:
            logger.info("No user info found in token data")
            return JSONResponse({
                "authenticated": False,
                "message": "No user information available"
            })
        
        logger.info(f"User {user_info.get('name')} is authenticated")
        return JSONResponse({
            "authenticated": True,
            "user": user_info
        })
        
    except Exception as e:
        logger.error(f"Error checking authentication status: {str(e)}")
        return JSONResponse({
            "authenticated": False,
            "message": f"Error checking authentication status: {str(e)}"
        })

@router.post("/logout")
async def logout(currentUser: Dict[str, Any] = Depends(getCurrentActiveUser)):
    """Logout from Microsoft"""
    try:
        # Get current user context
        mandateId, userId = await getUserContext(currentUser)
        if not mandateId or not userId:
            return JSONResponse({
                "message": "Not authenticated with Microsoft"
            })
            
        # Get LucyDOM interface for current user
        mydom = getLucydomInterface(
            mandateId=mandateId,
            userId=userId
        )
        if not mydom:
            return JSONResponse({
                "message": "Not authenticated with Microsoft"
            })
            
        # Remove token from database
        tokens = mydom.db.getRecordset("msftTokens", recordFilter={
            "mandateId": mandateId,
            "userId": userId
        })
        
        if tokens and len(tokens) > 0:
            mydom.db.recordDelete("msftTokens", tokens[0]["id"])
            logger.info(f"Removed Microsoft token for user {userId}")
        
        return JSONResponse({
            "message": "Successfully logged out from Microsoft"
        })
        
    except Exception as e:
        logger.error(f"Error during logout: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Logout failed: {str(e)}"
        )

@router.get("/token")
async def get_access_token(currentUser: Dict[str, Any] = Depends(getCurrentActiveUser)):
    """Get the current user's access token for Microsoft Graph API"""
    try:
        # Check if we have a token for the current user
        token_data = await load_token_from_file(currentUser)
        
        if not token_data:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Not authenticated with Microsoft"
            )
        
        # Verify token is still valid
        if not verify_token(token_data["access_token"]):
            # Try to refresh the token
            if not await refresh_token(currentUser["id"], currentUser):
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="Token expired and refresh failed"
                )
            # Reload token data after refresh
            token_data = await load_token_from_file(currentUser)
        
        return JSONResponse({
            "access_token": token_data["access_token"]
        })
        
    except Exception as e:
        logger.error(f"Error getting access token: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Error getting access token: {str(e)}"
        )

@router.post("/token")
async def get_backend_token(request: Request):
    """Convert MSAL token to backend token"""
    try:
        # Get the authorization header
        auth_header = request.headers.get('Authorization')
        if not auth_header or not auth_header.startswith('Bearer '):
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Missing or invalid authorization header"
            )

        # Extract the MSAL token
        msal_token = auth_header.split(' ')[1]

        # Verify the MSAL token and get user info
        user_info = get_user_info_from_token(msal_token)
        if not user_info:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Invalid MSAL token"
            )

        # Get the user from the database using the email
        gateway = getGatewayInterface()
        user = gateway.getUserByUsername(user_info["email"])
        
        if not user:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="User not registered in the system"
            )

        # Create backend token
        access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
        access_token = createAccessToken(
            data={
                "sub": user["username"],
                "mandateId": user["mandateId"]
            },
            expiresDelta=access_token_expires
        )

        return {
            "accessToken": access_token,
            "tokenType": "bearer",
            "user": {
                "username": user["username"],
                "email": user["email"],
                "fullName": user.get("fullName", ""),
                "mandateId": user["mandateId"]
            }
        }

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error in MSAL token conversion: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Error processing MSAL token: {str(e)}"
        )

@router.post("/save-token")
async def save_token(token_data: Dict[str, Any], currentUser: Dict[str, Any] = Depends(getCurrentActiveUser)):
    """Save Microsoft token data from frontend"""
    try:
        # Save token to database
        success = await save_token_to_file(token_data, currentUser)
        
        if success:
            return JSONResponse({
                "success": True,
                "message": "Token saved successfully"
            })
        else:
            return JSONResponse({
                "success": False,
                "message": "Failed to save token"
            })
            
    except Exception as e:
        logger.error(f"Error saving token: {str(e)}")
        return JSONResponse({
            "success": False,
            "message": f"Error saving token: {str(e)}"
        })