sana0721/vertex

1from fastapi import HTTPException, Header, Depends2from fastapi.security import APIKeyHeader3from typing import Optional4from config import API_KEY, HUGGINGFACE_API_KEY, HUGGINGFACE # Import API_KEY, HUGGINGFACE_API_KEY, HUGGINGFACE5import os6import json7import base648 9# Function to validate API key (moved from config.py)10def validate_api_key(api_key_to_validate: str) -> bool:11    """12    Validate the provided API key against the configured key.13    """14    if not API_KEY: # API_KEY is imported from config15        # If no API key is configured, authentication is disabled (or treat as invalid)16        # Depending on desired behavior, for now, let's assume if API_KEY is not set, all keys are invalid unless it's an empty string match17        return False # Or True if you want to disable auth when API_KEY is not set18    return api_key_to_validate == API_KEY19 20# API Key security scheme21api_key_header = APIKeyHeader(name="Authorization", auto_error=False)22 23# Dependency for API key validation24async def get_api_key(25    authorization: Optional[str] = Header(None),26    x_ip_token: Optional[str] = Header(None, alias="x-ip-token")27):28    # Check if Hugging Face auth is enabled29    if HUGGINGFACE:  # Use HUGGINGFACE from config30        if x_ip_token is None:31            raise HTTPException(32                status_code=401, # Unauthorised - because x-ip-token is missing33                detail="Missing x-ip-token header. This header is required for Hugging Face authentication."34            )35 36        try:37            # Decode JWT payload38            parts = x_ip_token.split('.')39            if len(parts) < 2:40                raise ValueError("Invalid JWT format: Not enough parts to extract payload.")41            payload_encoded = parts[1]42            # Add padding if necessary, as Python's base64.urlsafe_b64decode requires it43            payload_encoded += '=' * (-len(payload_encoded) % 4)44            decoded_payload_bytes = base64.urlsafe_b64decode(payload_encoded)45            payload = json.loads(decoded_payload_bytes.decode('utf-8'))46        except ValueError as ve:47            # Log server-side for debugging, but return a generic client error48            print(f"ValueError processing x-ip-token: {ve}")49            raise HTTPException(status_code=400, detail=f"Invalid JWT format in x-ip-token: {str(ve)}")50        except (json.JSONDecodeError, base64.binascii.Error, UnicodeDecodeError) as e:51            print(f"Error decoding/parsing x-ip-token payload: {e}")52            raise HTTPException(status_code=400, detail=f"Malformed x-ip-token payload: {str(e)}")53        except Exception as e: # Catch any other unexpected errors during token processing54            print(f"Unexpected error processing x-ip-token: {e}")55            raise HTTPException(status_code=500, detail="Internal error processing x-ip-token.")56 57        error_in_token = payload.get("error")58 59        if error_in_token == "InvalidAccessToken":60            raise HTTPException(61                status_code=403,62                detail="Access denied: x-ip-token indicates 'InvalidAccessToken'."63            )64        elif error_in_token is None:  # JSON 'null' is Python's None65            # If error is null, auth is successful. Now check if HUGGINGFACE_API_KEY is configured.66            print(f"HuggingFace authentication successful via x-ip-token (error field was null).")67            return HUGGINGFACE_API_KEY # Return the configured HUGGINGFACE_API_KEY68        else:69            # Any other non-null, non-"InvalidAccessToken" value in 'error' field70            raise HTTPException(71                status_code=403,72                detail=f"Access denied: x-ip-token indicates an unhandled error: '{error_in_token}'."73            )74    else:75        # Fallback to Bearer token authentication if HUGGINGFACE env var is not "true"76        if authorization is None:77            detail_message = "Missing API key. Please include 'Authorization: Bearer YOUR_API_KEY' header."78            # Optionally, provide a hint if the HUGGINGFACE env var exists but is not "true"79            if os.getenv("HUGGINGFACE") is not None: # Check for existence, not value80                 detail_message += " (Note: HUGGINGFACE mode with x-ip-token is not currently active)."81            raise HTTPException(82                status_code=401,83                detail=detail_message84            )85        86        # Check if the header starts with "Bearer "87        if not authorization.startswith("Bearer "):88            raise HTTPException(89                status_code=401,90                detail="Invalid API key format. Use 'Authorization: Bearer YOUR_API_KEY'"91            )92        93        # Extract the API key94        api_key = authorization.replace("Bearer ", "")95        96        # Validate the API key97        if not validate_api_key(api_key): # Call local validate_api_key98            raise HTTPException(99                status_code=401,100                detail="Invalid API key"101            )102        103        return api_key