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
        auth_url, _ = flow.authorization_url(
            access_type="offline",
            include_granted_scopes="true",
            state=json.dumps({
                "type": state,
                "connectionId": connectionId
            })
        )
        
        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")
        
        # 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,
                authority=AuthAuthority.GOOGLE,
                tokenAccess=credentials.token,
                tokenRefresh=credentials.refresh_token,
                tokenType=credentials.token_type,
                expiresAt=credentials.expiry.timestamp() if credentials.expiry else None
            )
            
            # 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:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="Connection ID is required for connection flow"
                )
            
            # Get current user from session
            current_user = await getCurrentUser(request)
            if not current_user:
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="User not authenticated"
                )
            
            # Find and update connection
            interface = getInterface(current_user)
            connection = None
            for conn in current_user.connections:
                if conn.id == connection_id:
                    connection = conn
                    break
            
            if not connection:
                raise HTTPException(
                    status_code=status.HTTP_404_NOT_FOUND,
                    detail="Connection not found"
                )
            
            # Update connection
            connection.status = ConnectionStatus.ACTIVE
            connection.lastChecked = datetime.now()
            connection.expiresAt = credentials.expiry if credentials.expiry else None
            
            # Update user record
            interface.db.recordModify("users", current_user.id, {
                "connections": [c.to_dict() for c in current_user.connections]
            })
            
            # Return success page
            return HTMLResponse(
                content=f"""
                <html>
                    <head><title>Connection Successful</title></head>
                    <body>
                        <script>
                            if (window.opener) {{
                                window.opener.postMessage({{
                                    type: 'google_connection_success',
                                    connectionId: {json.dumps(connection_id)}
                                }}, '*');
                            }}
                            setTimeout(() => window.close(), 1000);
                        </script>
                    </body>
                </html>
                """
            )
        
    except Exception as e:
        logger.error(f"Error in auth callback: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Authentication failed: {str(e)}"
        )

@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)}"
        )