security Updated

app.py

@@ -1,5 +1,6 @@
-from fastapi import FastAPI, HTTPException, Request, UploadFile, File
+from fastapi import FastAPI, HTTPException, Request, UploadFile, File, Depends, status
 from fastapi.responses import StreamingResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel, Field
 from typing import List, Optional, Dict, Any, AsyncGenerator
 import asyncio
@@ -11,6 +12,9 @@ from contextlib import asynccontextmanager
 import tempfile
 import shutil
 import random
+import hashlib
+import secrets
+from functools import wraps
 
 # Third-party imports
 from openai import OpenAI, AsyncOpenAI
@@ -75,11 +79,170 @@ class Config:
     TOP_K = int(os.getenv("TOP_K", "10"))
     SIMILARITY_THRESHOLD = float(os.getenv("SIMILARITY_THRESHOLD", "0.1"))
     DEVICE = os.getenv("DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Security Configuration
+    API_KEYS = os.getenv("API_KEYS", "").split(",") if os.getenv("API_KEYS") else []
+    MASTER_KEY = os.getenv("MASTER_KEY", "")
+    ENABLE_SECURITY = os.getenv("ENABLE_SECURITY", "true").lower() == "true"
+    RATE_LIMIT_PER_MINUTE = int(os.getenv("RATE_LIMIT_PER_MINUTE", "60"))
+    
+    @classmethod
+    def generate_api_key(cls) -> str:
+        """Generate a new API key"""
+        return f"sk-{secrets.token_urlsafe(32)}"
+    
+    @classmethod
+    def validate_api_key(cls, api_key: str) -> bool:
+        """Validate API key"""
+        if not cls.ENABLE_SECURITY:
+            return True
+        
+        if not api_key:
+            return False
+        
+        # Check master key
+        if cls.MASTER_KEY and api_key == cls.MASTER_KEY:
+            return True
+        
+        # Check configured API keys
+        if cls.API_KEYS and api_key in cls.API_KEYS:
+            return True
+        
+        return False
 
 # Models
 OPENROUTER_MODELS = ["deepseek/deepseek-chat-v3-0324:free", "deepseek/deepseek-r1-0528:free", "qwen/qwen3-235b-a22b:free", "google/gemini-2.0-flash-exp:free"]
 GROQ_MODELS = ["llama-3.3-70b-versatile"]
 
+# Security Models
+class APIKeyRequest(BaseModel):
+    description: Optional[str] = Field(None, description="Description for the API key")
+
+class APIKeyResponse(BaseModel):
+    api_key: str
+    description: Optional[str] = None
+    created_at: str
+    status: str = "active"
+
+class SecurityInfo(BaseModel):
+    security_enabled: bool
+    rate_limit_per_minute: int
+    has_master_key: bool
+    configured_keys_count: int
+
+# Rate Limiting
+class RateLimiter:
+    def __init__(self):
+        self.requests = {}
+        self.blocked_ips = set()
+    
+    def is_allowed(self, identifier: str, limit_per_minute: int = Config.RATE_LIMIT_PER_MINUTE) -> bool:
+        """Check if request is allowed based on rate limit"""
+        if not Config.ENABLE_SECURITY:
+            return True
+        
+        if identifier in self.blocked_ips:
+            return False
+        
+        now = datetime.now()
+        minute_key = now.strftime("%Y-%m-%d %H:%M")
+        
+        if identifier not in self.requests:
+            self.requests[identifier] = {}
+        
+        if minute_key not in self.requests[identifier]:
+            self.requests[identifier][minute_key] = 0
+        
+        # Clean old entries (keep only last 2 minutes)
+        keys_to_remove = []
+        for key in self.requests[identifier]:
+            try:
+                key_time = datetime.strptime(key, "%Y-%m-%d %H:%M")
+                if (now - key_time).total_seconds() > 120:  # 2 minutes
+                    keys_to_remove.append(key)
+            except ValueError:
+                keys_to_remove.append(key)
+        
+        for key in keys_to_remove:
+            del self.requests[identifier][key]
+        
+        # Check current minute limit
+        current_requests = self.requests[identifier].get(minute_key, 0)
+        if current_requests >= limit_per_minute:
+            return False
+        
+        self.requests[identifier][minute_key] = current_requests + 1
+        return True
+    
+    def block_ip(self, ip: str):
+        """Block an IP address"""
+        self.blocked_ips.add(ip)
+    
+    def unblock_ip(self, ip: str):
+        """Unblock an IP address"""
+        self.blocked_ips.discard(ip)
+
+# Security Dependencies
+security = HTTPBearer(auto_error=False)
+rate_limiter = RateLimiter()
+
+async def verify_api_key(
+    request: Request,
+    credentials: Optional[HTTPAuthorizationCredentials] = Depends(security)
+) -> str:
+    """Verify API key from Authorization header"""
+    if not Config.ENABLE_SECURITY:
+        return "security_disabled"
+    
+    # Get client IP
+    client_ip = request.client.host
+    
+    # Check rate limit
+    if not rate_limiter.is_allowed(client_ip):
+        raise HTTPException(
+            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+            detail="Rate limit exceeded"
+        )
+    
+    # Check API key
+    if not credentials:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="API key required. Please provide a valid API key in the Authorization header as 'Bearer <your-api-key>'"
+        )
+    
+    api_key = credentials.credentials
+    if not Config.validate_api_key(api_key):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid API key"
+        )
+    
+    return api_key
+
+async def verify_master_key(
+    request: Request,
+    credentials: Optional[HTTPAuthorizationCredentials] = Depends(security)
+) -> str:
+    """Verify master key for admin operations"""
+    if not Config.ENABLE_SECURITY:
+        return "security_disabled"
+    
+    if not credentials:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Master key required for admin operations"
+        )
+    
+    api_key = credentials.credentials
+    if not Config.MASTER_KEY or api_key != Config.MASTER_KEY:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Invalid master key"
+        )
+    
+    return api_key
+
 class DynamicOpenAIService:
     """Service for dynamic OpenAI provider selection"""
     
@@ -656,6 +819,19 @@ async def lifespan(app: FastAPI):
     
     print("🚀 All services initialized successfully!")
     
+    # Print security information
+    if Config.ENABLE_SECURITY:
+        print("\n🔒 Security Configuration:")
+        print(f"  Security: ENABLED")
+        print(f"  Rate Limit: {Config.RATE_LIMIT_PER_MINUTE} requests/minute")
+        print(f"  Master Key: {'✓ Configured' if Config.MASTER_KEY else '✗ Not configured'}")
+        print(f"  API Keys: {len([k for k in Config.API_KEYS if k.strip()])} configured")
+        if not Config.MASTER_KEY and not Config.API_KEYS:
+            print("  ⚠️  WARNING: No API keys configured! Set MASTER_KEY or API_KEYS environment variable.")
+    else:
+        print("\n🔓 Security: DISABLED")
+        print("  All endpoints are publicly accessible")
+    
     yield
     
     # Shutdown
@@ -681,10 +857,15 @@ app = FastAPI(
 
 @app.get("/")
 async def root():
-    return {"message": "Enhanced RAG API with Dynamic Provider Selection", "status": "running"}
+    return {
+        "message": "Enhanced RAG API with Dynamic Provider Selection", 
+        "status": "running",
+        "security_enabled": Config.ENABLE_SECURITY,
+        "version": "1.0.0"
+    }
 
 @app.get("/health")
-async def health_check():
+async def health_check(api_key: str = Depends(verify_api_key)):
     """Health check endpoint"""
     try:
         # Test Qdrant connection
@@ -755,7 +936,7 @@ async def health_check():
     }
 
 @app.post("/v1/chat/completions")
-async def chat_completions(request: ChatCompletionRequest):
+async def chat_completions(request: ChatCompletionRequest, api_key: str = Depends(verify_api_key)):
     """OpenAI-compatible chat completions endpoint with enhanced RAG and dynamic provider selection"""
     
     if not app_state.openai_service:
@@ -912,7 +1093,11 @@ async def stream_chat_completion(messages: List[Dict], request: ChatCompletionRe
 
 # Document management endpoints
 @app.post("/v1/documents/upload")
-async def upload_document(file: UploadFile = File(...), metadata: str = None):
+async def upload_document(
+    file: UploadFile = File(...), 
+    metadata: str = None,
+    api_key: str = Depends(verify_api_key)
+):
     """Upload a PDF document"""
     try:
         if not app_state.document_manager:
@@ -966,7 +1151,7 @@ async def upload_document(file: UploadFile = File(...), metadata: str = None):
         raise HTTPException(status_code=500, detail=f"Error uploading document: {str(e)}")
 
 @app.post("/v1/documents/search")
-async def search_documents(request: DocumentSearchRequest):
+async def search_documents(request: DocumentSearchRequest, api_key: str = Depends(verify_api_key)):
     """Search for documents"""
     try:
         if not app_state.document_manager:
@@ -989,7 +1174,7 @@ async def search_documents(request: DocumentSearchRequest):
         raise HTTPException(status_code=500, detail=f"Error searching documents: {str(e)}")
 
 @app.get("/v1/documents/list")
-async def list_documents():
+async def list_documents(api_key: str = Depends(verify_api_key)):
     """List all documents"""
     try:
         if not app_state.document_manager:
@@ -1007,7 +1192,7 @@ async def list_documents():
         raise HTTPException(status_code=500, detail=f"Error listing documents: {str(e)}")
 
 @app.delete("/v1/documents/{document_id}")
-async def delete_document(document_id: str):
+async def delete_document(document_id: str, api_key: str = Depends(verify_api_key)):
     """Delete a document"""
     try:
         if not app_state.document_manager:
@@ -1028,7 +1213,7 @@ async def delete_document(document_id: str):
 
 # Legacy compatibility endpoints
 @app.post("/v1/embeddings/add")
-async def add_document_legacy(content: str, metadata: Optional[Dict] = None):
+async def add_document_legacy(content: str, metadata: Optional[Dict] = None, api_key: str = Depends(verify_api_key)):
     """Legacy endpoint for adding documents (text content)"""
     try:
         if not app_state.embedding_service or not app_state.qdrant_client:
@@ -1059,7 +1244,7 @@ async def add_document_legacy(content: str, metadata: Optional[Dict] = None):
         raise HTTPException(status_code=500, detail=f"Error adding document: {str(e)}")
 
 @app.get("/v1/collections/info")
-async def get_collection_info():
+async def get_collection_info(api_key: str = Depends(verify_api_key)):
     """Get information about the collection"""
     try:
         if app_state.qdrant_client is None:
@@ -1079,7 +1264,7 @@ async def get_collection_info():
 
 # New endpoint to get available providers and models
 @app.get("/v1/providers")
-async def get_providers():
+async def get_providers(api_key: str = Depends(verify_api_key)):
     """Get available providers and their models"""
     try:
         if not app_state.openai_service:
@@ -1117,6 +1302,100 @@ async def get_providers():
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error getting providers: {str(e)}")
 
+# Security Management Endpoints
+@app.get("/v1/security/info")
+async def get_security_info() -> SecurityInfo:
+    """Get security configuration information (public endpoint)"""
+    return SecurityInfo(
+        security_enabled=Config.ENABLE_SECURITY,
+        rate_limit_per_minute=Config.RATE_LIMIT_PER_MINUTE,
+        has_master_key=bool(Config.MASTER_KEY),
+        configured_keys_count=len([k for k in Config.API_KEYS if k.strip()])
+    )
+
+@app.post("/v1/security/generate-key")
+async def generate_api_key(
+    request: APIKeyRequest,
+    master_key: str = Depends(verify_master_key)
+) -> APIKeyResponse:
+    """Generate a new API key (requires master key)"""
+    try:
+        new_key = Config.generate_api_key()
+        
+        return APIKeyResponse(
+            api_key=new_key,
+            description=request.description,
+            created_at=datetime.now().isoformat(),
+            status="active"
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error generating API key: {str(e)}")
+
+@app.post("/v1/security/validate-key")
+async def validate_api_key_endpoint(
+    api_key: str = Depends(verify_api_key)
+) -> Dict[str, Any]:
+    """Validate an API key"""
+    return {
+        "valid": True,
+        "key_type": "master" if api_key == Config.MASTER_KEY else "standard",
+        "validated_at": datetime.now().isoformat()
+    }
+
+@app.get("/v1/security/rate-limit-status")
+async def get_rate_limit_status(
+    request: Request,
+    api_key: str = Depends(verify_api_key)
+) -> Dict[str, Any]:
+    """Get current rate limit status"""
+    client_ip = request.client.host
+    
+    # Get current minute requests
+    now = datetime.now()
+    minute_key = now.strftime("%Y-%m-%d %H:%M")
+    
+    current_requests = 0
+    if client_ip in rate_limiter.requests:
+        current_requests = rate_limiter.requests[client_ip].get(minute_key, 0)
+    
+    return {
+        "client_ip": client_ip,
+        "current_requests": current_requests,
+        "limit_per_minute": Config.RATE_LIMIT_PER_MINUTE,
+        "remaining_requests": max(0, Config.RATE_LIMIT_PER_MINUTE - current_requests),
+        "reset_at": f"{minute_key}:00",
+        "is_blocked": client_ip in rate_limiter.blocked_ips
+    }
+
+# Admin endpoints for IP management
+@app.post("/v1/admin/block-ip/{ip}")
+async def block_ip(
+    ip: str,
+    master_key: str = Depends(verify_master_key)
+) -> Dict[str, str]:
+    """Block an IP address (requires master key)"""
+    rate_limiter.block_ip(ip)
+    return {"message": f"IP {ip} has been blocked", "blocked_at": datetime.now().isoformat()}
+
+@app.post("/v1/admin/unblock-ip/{ip}")
+async def unblock_ip(
+    ip: str,
+    master_key: str = Depends(verify_master_key)
+) -> Dict[str, str]:
+    """Unblock an IP address (requires master key)"""
+    rate_limiter.unblock_ip(ip)
+    return {"message": f"IP {ip} has been unblocked", "unblocked_at": datetime.now().isoformat()}
+
+@app.get("/v1/admin/blocked-ips")
+async def get_blocked_ips(
+    master_key: str = Depends(verify_master_key)
+) -> Dict[str, Any]:
+    """Get list of blocked IPs (requires master key)"""
+    return {
+        "blocked_ips": list(rate_limiter.blocked_ips),
+        "count": len(rate_limiter.blocked_ips)
+    }
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)