Add API key auth and admin management routes

.gitignore

@@ -4,3 +4,6 @@ __pycache__/
 telegram_bot.py
 bot.py
 render_bot/
+web/
+api_keys.json
+PROMPT_FOR_OPUS_VERCEL.md

auth.py

@@ -0,0 +1,164 @@
+import os
+import json
+import time
+import secrets
+import threading
+from typing import Optional, Dict, Any, List, Tuple
+
+KEYS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "api_keys.json")
+ADMIN_MASTER_KEY = os.environ.get("ADMIN_MASTER_KEY", "secretz_master_admin_2026")
+INTERNAL_BOT_KEY = os.environ.get("INTERNAL_BOT_KEY", "sk_internal_telebot_secretz_9921")
+REQUIRE_AUTH = os.environ.get("REQUIRE_AUTH", "false").lower() in ("true", "1", "yes")
+
+_lock = threading.Lock()
+# In-memory sliding window for rate limiting: {api_key: [timestamp1, timestamp2, ...]}
+_rate_limits: Dict[str, List[float]] = {}
+
+
+def load_keys() -> Dict[str, Dict[str, Any]]:
+    if os.path.exists(KEYS_FILE):
+        try:
+            with open(KEYS_FILE, "r", encoding="utf-8") as f:
+                return json.load(f)
+        except Exception:
+            return {}
+    return {}
+
+
+def save_keys(keys_data: Dict[str, Dict[str, Any]]):
+    try:
+        with open(KEYS_FILE, "w", encoding="utf-8") as f:
+            json.dump(keys_data, f, indent=2)
+    except Exception as e:
+        print(f"Error saving api_keys.json: {e}")
+
+
+class KeyManager:
+    @staticmethod
+    def create_key(
+        name: str,
+        days_valid: Optional[int] = 30,
+        rate_limit_per_min: int = 60,
+        total_credits: int = 1000
+    ) -> Dict[str, Any]:
+        """Creates a new API key with expiration and credit limits."""
+        api_key = f"sk_live_{secrets.token_hex(16)}"
+        now = time.time()
+        expires_at = now + (days_valid * 86400) if (days_valid and days_valid > 0) else None
+
+        key_info = {
+            "api_key": api_key,
+            "name": name.strip() or "Unnamed Client",
+            "created_at": now,
+            "expires_at": expires_at,
+            "rate_limit_per_min": max(1, int(rate_limit_per_min)),
+            "total_credits": int(total_credits),  # -1 for unlimited
+            "used_credits": 0,
+            "status": "active"
+        }
+
+        with _lock:
+            keys = load_keys()
+            keys[api_key] = key_info
+            save_keys(keys)
+
+        return key_info
+
+    @staticmethod
+    def list_keys() -> List[Dict[str, Any]]:
+        with _lock:
+            keys = load_keys()
+            return list(keys.values())
+
+    @staticmethod
+    def revoke_key(api_key: str) -> bool:
+        with _lock:
+            keys = load_keys()
+            if api_key in keys:
+                keys[api_key]["status"] = "revoked"
+                save_keys(keys)
+                return True
+            return False
+
+    @staticmethod
+    def delete_key(api_key: str) -> bool:
+        with _lock:
+            keys = load_keys()
+            if api_key in keys:
+                del keys[api_key]
+                save_keys(keys)
+                return True
+            return False
+
+    @staticmethod
+    def validate_and_consume(api_key: Optional[str]) -> Tuple[bool, int, str]:
+        """
+        Validates API key, checks expiration, rate limit, and consumes 1 credit.
+        Returns: (is_valid: bool, error_code: int, error_message: str)
+        """
+        if not REQUIRE_AUTH:
+            return True, 200, "OK"
+
+        if not api_key:
+            return False, 401, "API Key required. Pass header 'X-API-Key' or query parameter 'api_key'."
+
+        # Master admin key bypass
+        if api_key == ADMIN_MASTER_KEY or api_key == INTERNAL_BOT_KEY:
+            return True, 200, "OK"
+
+        with _lock:
+            keys = load_keys()
+            if api_key not in keys:
+                return False, 401, "Invalid API Key."
+
+            info = keys[api_key]
+
+            # 1. Status check
+            if info.get("status") != "active":
+                return False, 403, "API Key has been revoked or deactivated."
+
+            now = time.time()
+
+            # 2. Expiry check
+            expires_at = info.get("expires_at")
+            if expires_at and now > expires_at:
+                info["status"] = "expired"
+                save_keys(keys)
+                return False, 403, "API Key has expired."
+
+            # 3. Credit check
+            total_credits = info.get("total_credits", -1)
+            used_credits = info.get("used_credits", 0)
+            if total_credits != -1 and used_credits >= total_credits:
+                return False, 403, "API Key credit quota exhausted."
+
+            # 4. Rate limit check (sliding window 60s)
+            rate_limit = info.get("rate_limit_per_min", 60)
+            if api_key not in _rate_limits:
+                _rate_limits[api_key] = []
+
+            # Keep only timestamps in the last 60 seconds
+            cutoff = now - 60.0
+            _rate_limits[api_key] = [t for t in _rate_limits[api_key] if t > cutoff]
+
+            if len(_rate_limits[api_key]) >= rate_limit:
+                retry_after = round(60.0 - (now - _rate_limits[api_key][0]), 1)
+                return False, 429, f"Rate limit exceeded ({rate_limit} req/min). Retry after {retry_after}s."
+
+            # Record request & consume credit
+            _rate_limits[api_key].append(now)
+            info["used_credits"] = used_credits + 1
+            save_keys(keys)
+
+            return True, 200, "OK"
+
+    @staticmethod
+    def get_key_status(api_key: str) -> Optional[Dict[str, Any]]:
+        if api_key == ADMIN_MASTER_KEY:
+            return {"name": "Master Admin", "status": "active", "total_credits": -1, "used_credits": 0, "expires_at": None}
+        if api_key == INTERNAL_BOT_KEY:
+            return {"name": "Telegram Bot Internal", "status": "active", "total_credits": -1, "used_credits": 0, "expires_at": None}
+
+        with _lock:
+            keys = load_keys()
+            return keys.get(api_key)

main.py

@@ -1,15 +1,17 @@
 import time
 import os
 from contextlib import asynccontextmanager
-from typing import Optional
+from typing import Optional, Dict, Any, List
 
-from fastapi import FastAPI, Query, HTTPException, Request
+from fastapi import FastAPI, Query, HTTPException, Request, Header, Depends
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
 from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel, Field
 
 from engine import LookupEngine
+from auth import KeyManager, ADMIN_MASTER_KEY, INTERNAL_BOT_KEY
 
 start_time = time.time()
 _engine: Optional[LookupEngine] = None
@@ -27,13 +29,12 @@ def get_engine() -> LookupEngine:
 async def lifespan(app: FastAPI):
     get_engine()
     yield
-    # Cleanup if needed
 
 
 app = FastAPI(
-    title="HyperLookup API",
-    description="Enterprise-grade ultra-fast lookup API for 2.5B+ Indian Records (Phone & Aadhar) hosted with $0 cost.",
-    version="2.0.0",
+    title="HyperLookup API & SaaS Engine",
+    description="Enterprise-grade ultra-fast lookup API for 2.5B+ Indian Records (Phone & Aadhar) with API Key Auth & Quota Controls.",
+    version="2.1.0",
     lifespan=lifespan,
     docs_url="/docs",
     redoc_url="/redoc"
@@ -42,7 +43,7 @@ app = FastAPI(
 # Enable GZip compression
 app.add_middleware(GZipMiddleware, minimum_size=1000)
 
-# Enable CORS for any frontend to consume this API freely
+# Enable CORS for Vercel, Cloudflare, Localhost and all frontend apps
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -51,33 +52,106 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Mount static files
+# Mount static directory if exists
 static_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")
 if os.path.exists(static_dir):
     app.mount("/static", StaticFiles(directory=static_dir), name="static")
 
 
+# ---------------------------------------------------------
+# Request Models
+# ---------------------------------------------------------
+class CreateKeyRequest(BaseModel):
+    name: str = Field(..., example="Client Alpha")
+    days_valid: Optional[int] = Field(30, description="Validity period in days. Pass 0 or null for lifetime.")
+    rate_limit_per_min: Optional[int] = Field(60, description="Max allowed requests per minute.")
+    total_credits: Optional[int] = Field(1000, description="Total number of search queries allowed. Pass -1 for unlimited.")
+
+
+# ---------------------------------------------------------
+# Authentication Dependency
+# ---------------------------------------------------------
+def verify_api_key(
+    request: Request,
+    api_key_query: Optional[str] = Query(None, alias="api_key"),
+    x_api_key: Optional[str] = Header(None, alias="X-API-Key"),
+    authorization: Optional[str] = Header(None)
+) -> str:
+    """Extracts and verifies API key from Header or Query Param."""
+    token = None
+    if x_api_key:
+        token = x_api_key
+    elif authorization and authorization.lower().startswith("bearer "):
+        token = authorization.split(" ", 1)[1]
+    elif api_key_query:
+        token = api_key_query
+
+    is_valid, code, msg = KeyManager.validate_and_consume(token)
+    if not is_valid:
+        raise HTTPException(status_code=code, detail=msg)
+
+    return token or "unauthenticated"
+
+
+def verify_admin_key(
+    x_admin_key: Optional[str] = Header(None, alias="X-Admin-Key"),
+    admin_key_query: Optional[str] = Query(None, alias="admin_key")
+):
+    """Requires the master admin key to access administrative operations."""
+    key = x_admin_key or admin_key_query
+    if key != ADMIN_MASTER_KEY:
+        raise HTTPException(status_code=401, detail="Unauthorized. Master Admin Key required.")
+    return key
+
+
+# ---------------------------------------------------------
+# Public Endpoints
+# ---------------------------------------------------------
 @app.get("/", response_class=HTMLResponse)
 async def serve_home():
-    """Serves the SaaS dashboard UI."""
     index_file = os.path.join(static_dir, "index.html")
     if os.path.exists(index_file):
         return FileResponse(index_file)
-    return HTMLResponse("<h1>HyperLookup API is Running</h1><p>Visit <a href='/docs'>/docs</a> for Swagger UI.</p>")
+    return HTMLResponse("<h1>HyperLookup Engine API is Online</h1><p>API documentation available at <a href='/docs'>/docs</a>.</p>")
+
 
+@app.get("/api/health")
+async def health_check():
+    """Liveness probe for monitoring services."""
+    return {"status": "ok", "timestamp": time.time()}
 
+
+@app.get("/api/stats")
+async def get_system_stats():
+    """Returns engine health, cache status, and database metrics."""
+    eng = get_engine()
+    uptime_sec = round(time.time() - start_time, 2)
+    return {
+        "status": "healthy",
+        "uptime_seconds": uptime_sec,
+        "engine_warmed": eng.is_warmed,
+        "cache_entries": len(eng.cache),
+        "database_info": {
+            "total_records_indexed": "2,500,000,000+",
+            "database_size": "227 GB (Compressed Columnar Parquet)",
+            "phone_partitions": 7,
+            "aadhar_partitions": 5,
+            "engine": "DuckDB Columnar Vectorized Engine",
+            "storage_layer": "Hugging Face Storage Buckets + Cloudflare CDN"
+        }
+    }
+
+
+# ---------------------------------------------------------
+# Protected Search Endpoints (Requires Valid API Key)
+# ---------------------------------------------------------
 @app.get("/api/lookup")
 async def unified_lookup(
-    q: Optional[str] = Query(None, description="Search query (phone number or aadhar number)"),
+    q: Optional[str] = Query(None, description="Search query (phone or aadhar)"),
     phone: Optional[str] = Query(None, description="Direct 10-digit phone number"),
-    aadhar: Optional[str] = Query(None, description="Direct 12-digit aadhar number")
+    aadhar: Optional[str] = Query(None, description="Direct 12-digit aadhar number"),
+    _key: str = Depends(verify_api_key)
 ):
-    """
-    Unified Lookup Endpoint:
-    - Automatically identifies whether input is a Phone Number or Aadhar Number.
-    - Routes directly to the exact indexed Parquet partition.
-    - Returns deduplicated and sanitized records with execution latency in milliseconds.
-    """
     eng = get_engine()
     if phone:
         return eng.query_by_phone(phone)
@@ -90,43 +164,58 @@ async def unified_lookup(
 
 
 @app.get("/api/phone/{phone_number}")
-async def lookup_by_phone(phone_number: str):
-    """Direct REST endpoint for phone lookup."""
+async def lookup_by_phone(phone_number: str, _key: str = Depends(verify_api_key)):
     return get_engine().query_by_phone(phone_number)
 
 
 @app.get("/api/aadhar/{aadhar_number}")
-async def lookup_by_aadhar(aadhar_number: str):
-    """Direct REST endpoint for aadhar lookup."""
+async def lookup_by_aadhar(aadhar_number: str, _key: str = Depends(verify_api_key)):
     return get_engine().query_by_aadhar(aadhar_number)
 
 
-@app.get("/api/stats")
-async def get_system_stats():
-    """Returns database size, cache statistics, partition status, and engine uptime."""
-    eng = get_engine()
-    uptime_sec = round(time.time() - start_time, 2)
-    cache_items = len(eng.cache)
+@app.get("/api/key/status")
+async def check_key_status(api_key: str = Query(..., description="API Key to inspect")):
+    status = KeyManager.get_key_status(api_key)
+    if not status:
+        raise HTTPException(status_code=404, detail="API Key not found.")
+    return status
+
+
+# ---------------------------------------------------------
+# Admin Key Management Endpoints (Master Admin Only)
+# ---------------------------------------------------------
+@app.post("/api/admin/keys/create")
+async def create_api_key(req: CreateKeyRequest, _admin: str = Depends(verify_admin_key)):
+    """Generate a new client API key with custom expiry, credits, and rate limit."""
+    new_key = KeyManager.create_key(
+        name=req.name,
+        days_valid=req.days_valid,
+        rate_limit_per_min=req.rate_limit_per_min or 60,
+        total_credits=req.total_credits if req.total_credits is not None else 1000
+    )
     return {
-        "status": "healthy",
-        "uptime_seconds": uptime_sec,
-        "engine_warmed": eng.is_warmed,
-        "cache_entries": cache_items,
-        "database_info": {
-            "total_records_indexed": "2,500,000,000+",
-            "database_size": "227 GB (Compressed Columnar Parquet)",
-            "phone_partitions": 7,
-            "aadhar_partitions": 5,
-            "engine": "DuckDB Columnar Vectorized Engine",
-            "storage_layer": "Hugging Face Storage Buckets + Cloudflare CDN"
-        }
+        "success": True,
+        "message": "API key generated successfully.",
+        "key_data": new_key
     }
 
 
-@app.get("/api/health")
-async def health_check():
-    """Liveness probe for deployment platforms."""
-    return {"status": "ok", "timestamp": time.time()}
+@app.get("/api/admin/keys")
+async def list_all_keys(_admin: str = Depends(verify_admin_key)):
+    """List all issued API keys with remaining credits and usage stats."""
+    return {
+        "success": True,
+        "keys": KeyManager.list_keys()
+    }
+
+
+@app.delete("/api/admin/keys/{target_key}")
+async def delete_api_key(target_key: str, _admin: str = Depends(verify_admin_key)):
+    """Permanently delete or revoke an API key."""
+    deleted = KeyManager.delete_key(target_key)
+    if not deleted:
+        raise HTTPException(status_code=404, detail="Key not found.")
+    return {"success": True, "message": "API key deleted successfully."}
 
 
 if __name__ == "__main__":