Update app.py

app.py

@@ -2,11 +2,13 @@
 Surya OCR - Universal REST API
 One endpoint: POST /ocr — any file in, extracted text out.
 Supports: Images, PDF, Word (.docx/.doc), PowerPoint (.pptx/.ppt), and more.
+Secret key authentication via X-API-Key header.
 Hugging Face Spaces (Free Tier) compatible.
 """
 
 import io
 import gc
+import os
 import logging
 import subprocess
 import tempfile
@@ -14,25 +16,38 @@ from contextlib import asynccontextmanager
 from pathlib import Path
 
 import uvicorn
-from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi import FastAPI, File, UploadFile, HTTPException, Security
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
+from fastapi.security import APIKeyHeader
 from PIL import Image
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# ── Auth ──────────────────────────────────────────────────────────────────────
+# Set your secret key as an environment variable in HF Space Settings:
+#   Settings → Variables and secrets → New secret → Name: API_SECRET_KEY
+API_SECRET_KEY = os.environ.get("API_SECRET_KEY", "")
+
+api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
+
+def verify_key(api_key: str = Security(api_key_header)):
+    if not API_SECRET_KEY:
+        raise RuntimeError("API_SECRET_KEY environment variable is not set on the server.")
+    if api_key != API_SECRET_KEY:
+        raise HTTPException(status_code=401, detail="Invalid or missing API key. Pass it as X-API-Key header.")
+    return api_key
+
 # ── Global model holders ──────────────────────────────────────────────────────
 recognition_predictor = None
 detection_predictor = None
 
 MAX_FILE_SIZE_MB = 30
-MAX_PAGES = 10  # Free tier memory guard
+MAX_PAGES = 10
 
 SUPPORTED_EXTENSIONS = {
-    # Images
     ".jpg", ".jpeg", ".png", ".webp", ".tiff", ".tif", ".bmp",
-    # Documents
     ".pdf",
     ".docx", ".doc",
     ".pptx", ".ppt",
@@ -70,7 +85,7 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(
     title="Surya OCR API",
-    description="Upload any document — image, PDF, Word, PowerPoint — get back extracted text.",
+    description="Upload any document — image, PDF, Word, PowerPoint — get back extracted text. Requires X-API-Key header.",
     version="2.0.0",
     lifespan=lifespan,
 )
@@ -90,10 +105,6 @@ def convert_pdf(data: bytes) -> list:
 
 
 def convert_office(data: bytes, suffix: str) -> list:
-    """
-    Convert Word / PowerPoint / ODT / ODP / RTF to PDF via LibreOffice,
-    then rasterize with pdf2image.
-    """
     with tempfile.TemporaryDirectory() as tmpdir:
         input_path = Path(tmpdir) / f"input{suffix}"
         input_path.write_bytes(data)
@@ -118,7 +129,6 @@ def convert_office(data: bytes, suffix: str) -> list:
 
 
 def file_to_images(data: bytes, suffix: str) -> list:
-    """Route any file to the right converter and return a list of PIL Images."""
     image_exts = {".jpg", ".jpeg", ".png", ".webp", ".tiff", ".tif", ".bmp"}
     office_exts = {".docx", ".doc", ".pptx", ".ppt", ".odt", ".odp", ".rtf"}
 
@@ -155,31 +165,37 @@ def run_ocr(images: list) -> str:
 
 @app.get("/", summary="Health / Info")
 def root():
+    """Public endpoint — no key required."""
     return {
         "service": "Surya OCR API",
         "status": "ready" if recognition_predictor else "loading",
-        "usage": "POST /ocr  with any file attached as form-data field 'file'",
+        "usage": "POST /ocr with your file as form-data field 'file' and X-API-Key header",
         "supported_formats": sorted(SUPPORTED_EXTENSIONS),
     }
 
 
-@app.get("/health")
+@app.get("/health", summary="Health check")
 def health():
+    """Public endpoint — no key required."""
     return {"status": "ready" if recognition_predictor else "loading"}
 
 
 @app.post("/ocr", summary="Extract text from any document")
-async def ocr(file: UploadFile = File(..., description="Any image, PDF, Word, or PowerPoint file")):
+async def ocr(
+    file: UploadFile = File(..., description="Any image, PDF, Word, or PowerPoint file"),
+    _: str = Security(verify_key),   # ← auth happens here
+):
     """
     Upload **any** document and receive the extracted text.
 
+    **Requires** `X-API-Key` header with your secret key.
+
     Accepts: JPG, PNG, WEBP, TIFF, BMP, PDF, DOCX, DOC, PPTX, PPT, ODT, ODP, RTF
 
-    Returns: { "text": "..." }
+    Returns: `{ "text": "..." }`
     """
     data = await file.read()
 
-    # Size check
     size_mb = len(data) / (1024 * 1024)
     if size_mb > MAX_FILE_SIZE_MB:
         raise HTTPException(status_code=413, detail=f"File too large ({size_mb:.1f} MB). Max: {MAX_FILE_SIZE_MB} MB.")