gateway/modules/routes/routeSecurityGoogle.py

"""
Routes for Google 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
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import Flow
from google.auth.transport.requests import Request as GoogleRequest
from googleapiclient.discovery import build

from modules.shared.configuration import APP_CONFIG
from modules.interfaces.serviceAppClass import getInterface, getRootInterface
from modules.interfaces.serviceAppModel import AuthAuthority, User, Token, ConnectionStatus, UserConnection
from modules.security.auth import getCurrentUser, limiter
from modules.shared.attributeUtils import ModelMixin

# Configure logger
logger = logging.getLogger(__name__)

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

# Google OAuth configuration
CLIENT_ID = APP_CONFIG.get("Service_GOOGLE_CLIENT_ID")
CLIENT_SECRET = APP_CONFIG.get("Service_GOOGLE_CLIENT_SECRET")
REDIRECT_URI = APP_CONFIG.get("Service_GOOGLE_REDIRECT_URI")
SCOPES = [
    "https://www.googleapis.com/auth/gmail.readonly",
    "https://www.googleapis.com/auth/userinfo.profile",
    "https://www.googleapis.com/auth/userinfo.email"
]

@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 Google login"""
    try:
        # Create OAuth flow
        flow = Flow.from_client_config(
            {
                "web": {
                    "client_id": CLIENT_ID,
                    "client_secret": CLIENT_SECRET,
                    "auth_uri": "https://accounts.google.com/o/oauth2/auth",
                    "token_uri": "https://oauth2.googleapis.com/token",
                    "redirect_uris": [REDIRECT_URI]
                }
            },
            scopes=SCOPES
        )
        
        # 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
            })
        
        # Generate auth URL with state
        auth_url, _ = flow.authorization_url(
            access_type="offline",
            include_granted_scopes="true",
            state=state_param,
            prompt="select_account"  # Force account selection screen
        )
        
        return RedirectResponse(auth_url)
        
    except Exception as e:
        logger.error(f"Error initiating Google login: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to initiate Google login: {str(e)}"
        )

@router.get("/auth/callback")
async def auth_callback(code: str, state: str, request: Request) -> HTMLResponse:
    """Handle Google 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 Google auth callback: state_type={state_type}, connection_id={connection_id}, user_id={user_id}")
        
        # Create OAuth flow
        flow = Flow.from_client_config(
            {
                "web": {
                    "client_id": CLIENT_ID,
                    "client_secret": CLIENT_SECRET,
                    "auth_uri": "https://accounts.google.com/o/oauth2/auth",
                    "token_uri": "https://oauth2.googleapis.com/token",
                    "redirect_uris": [REDIRECT_URI]
                }
            },
            scopes=SCOPES
        )
        
        # Exchange code for credentials
        flow.fetch_token(code=code)
        credentials = flow.credentials
        
        # Get user info
        user_info_response = flow.oauth2session.get("https://www.googleapis.com/oauth2/v2/userinfo")
        user_info = user_info_response.json()
        
        if state_type == "login":
            # Handle login flow
            rootInterface = getRootInterface()
            user = rootInterface.getUserByUsername(user_info.get("email"))
            
            if not user:
                # Create new user if doesn't exist
                user = rootInterface.createUser(
                    username=user_info.get("email"),
                    email=user_info.get("email"),
                    fullName=user_info.get("name"),
                    authenticationAuthority=AuthAuthority.GOOGLE,
                    externalId=user_info.get("id"),
                    externalUsername=user_info.get("email"),
                    externalEmail=user_info.get("email")
                )
            
            # Create token
            token = Token(
                userId=user.id,  # Use local user's ID
                authority=AuthAuthority.GOOGLE,
                tokenAccess=credentials.token,
                tokenRefresh=credentials.refresh_token,
                tokenType=credentials.token_type,
                expiresAt=credentials.expiry.timestamp() if credentials.expiry else None,
                createdAt=datetime.now()
            )
            
            # Save token
            appInterface = getInterface(user)
            appInterface.saveToken(token)
            
            # 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: 'google_auth_success',
                                    access_token: {json.dumps(credentials.token)},
                                    token_data: {json.dumps(token.to_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: 'google_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: 'google_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 = datetime.now()
                connection.expiresAt = credentials.expiry if credentials.expiry else None
                connection.externalId = user_info.get("id")
                connection.externalUsername = user_info.get("email")
                connection.externalEmail = user_info.get("email")
                
                # Update connection record directly
                rootInterface.db.recordModify("connections", connection_id, connection.to_dict())
                
                # Save token
                token = Token(
                    userId=user.id,  # Use local user's ID
                    authority=AuthAuthority.GOOGLE,
                    tokenAccess=credentials.token,
                    tokenRefresh=credentials.refresh_token,
                    tokenType=credentials.token_type,
                    expiresAt=credentials.expiry.timestamp() if credentials.expiry else None,
                    createdAt=datetime.now()
                )
                interface.saveToken(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: 'google_connection_success',
                                        connection: {{
                                            id: '{connection.id}',
                                            status: 'connected',
                                            type: 'google',
                                            lastChecked: '{datetime.now().isoformat()}',
                                            expiresAt: '{credentials.expiry.isoformat() if credentials.expiry else None}'
                                        }}
                                    }}, '*');
                                    // 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: {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: 'google_connection_error',
                                        error: 'Failed to update connection: {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: 'google_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=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)}"
        )