Harden text and audio API for authenticated clients

app/api/analyze.py

@@ -1,120 +1,114 @@
-import os
 import asyncio
+import os
+import tempfile
+import threading
+
 import filetype
-from fastapi import APIRouter, HTTPException, WebSocketException, UploadFile, WebSocket, WebSocketDisconnect, Request
-from app.services import rate_limit
-from app.models.response import riskAssessment
-from app.models.request import information
-from app.services.risk import get_assessment
-from app.services.transcription import audio_transcript
-from app.config import MAX_FILE_SIZE
-import uuid
-router = APIRouter()
+from fastapi import APIRouter, Depends, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
 
-Allowed = {
-    "audio/mpeg",
-    "audio/m4a",
-    "audio/mp4",
-    "audio/wav",
-    "audio/x-wav",
-    "audio/webm",
-    "audio/ogg",
-    "audio/flac",
-}
-SUPPORTED_AUDIO_FORMATS = ("MP3", "M4A", "MP4", "WAV", "WebM", "OGG", "FLAC")
+from app.config import MAX_FILE_SIZE, CORS_ORIGINS
+from app.models.response import riskAssessment
+from app.models.request import TextRequest, information
+from app.services.auth import authorize, authorize_websocket, enforce_limit
+from app.services.risk import get_assessment, AssessmentUnavailable
+from app.services.transcription import audio_transcript, InvalidAudio
 
+router = APIRouter()
+# Serialize model inference per worker; reject excess work instead of accumulating uploads.
+audio_slot = threading.BoundedSemaphore(1)
+Allowed = {"audio/mpeg", "audio/m4a", "audio/mp4", "audio/wav", "audio/x-wav",
+           "audio/webm", "audio/ogg", "audio/flac"}
 UNSUPPORTED_AUDIO_ERROR = (
-    "Unsupported audio content type. "
-    f"Accepted formats: {', '.join(SUPPORTED_AUDIO_FORMATS)}."
+    "Unsupported audio content type. Accepted formats: MP3, M4A, MP4, WAV, WebM, OGG, FLAC."
 )
 
+
+def assess(text: str) -> riskAssessment:
+    try:
+        return get_assessment(text)
+    except AssessmentUnavailable as exc:
+        raise HTTPException(502, "Assessment temporarily unavailable") from exc
+
+
 @router.post(
-    "/email",
-    summary="Analyze email content for scam risk",
-    description=(
-        "Analyze the submitted email body and return a scam risk assessment. "
-        "Requests are rate-limited per client."
-    ),
-    responses={
-        429: {"description": "Rate limit exceeded."},
-    },
+    "/text", summary="Analyze text for scam risk", dependencies=[Depends(authorize)],
+    responses={401: {"description": "Authentication required"},
+               429: {"description": "Rate limit exceeded"}, 502: {"description": "Provider unavailable"}},
 )
-def email_check(item: information, request: Request)-> riskAssessment | dict | None:
-    assert request.client is not None
-    if not rate_limit.check_rate_limit(request.client.host):
-        raise HTTPException (status_code= 429, detail= {"error":"Reached your limit, wait 60 seconds before requesting again"})
-    return get_assessment(item.body)
+def text_check(item: TextRequest) -> riskAssessment:
+    return assess(item.body)
+
 
 @router.post(
-    "/audio",
-    summary="Analyze an audio file for scam risk",
-    description=(
-        "Upload a supported audio file for transcription and scam risk analysis. "
-        "The endpoint enforces the configured maximum file size and validates "
-        "the detected audio format."
-    ),
-    responses={
-        413: {"description": "Uploaded audio exceeds the configured size limit."},
-        415: {"description": "Uploaded content is not a supported audio format."},
-        429: {"description": "Rate limit exceeded."},
-    },
+    "/email", summary="Analyze email content for scam risk", deprecated=True,
+    dependencies=[Depends(authorize)],
+    responses={429: {"description": "Rate limit exceeded"}},
 )
-def audio_check(file:UploadFile, request: Request)-> riskAssessment | dict | None:
-
+def email_check(item: information) -> riskAssessment:
+    return assess(item.body)
 
-    assert request.client is not None
-    if not rate_limit.check_rate_limit(request.client.host):
-        raise HTTPException (status_code= 429, detail= {"error":"Reached your limit, wait 60 seconds before requesting again"})
-    byte =  file.file.read()
-    if len(byte) > MAX_FILE_SIZE:
-        raise HTTPException(status_code=413, detail={"error": f"File too large. Max size is {MAX_FILE_SIZE // (1024 * 1024)}MB."})
-    kind = filetype.guess(byte)
 
+def analyze_audio(data: bytes) -> riskAssessment:
+    if len(data) > MAX_FILE_SIZE:
+        raise HTTPException(413, {"error": f"File too large. Max size is {MAX_FILE_SIZE // (1024 * 1024)}MB."})
+    kind = filetype.guess(data)
     if kind is None or kind.mime not in Allowed:
-        raise HTTPException (status_code= 415, detail= {"error": UNSUPPORTED_AUDIO_ERROR})
-    tmp_dir = "/dev/shm/" if os.path.exists("/dev/shm") else ""
-    filename = f"{tmp_dir}audio{uuid.uuid4()}.{kind.extension}"
-    with open(filename, "wb") as f:
-        f.write(byte)
-    transcript = audio_transcript(filename)
-    os.remove(filename)
-    return get_assessment(transcript)
+        raise HTTPException(415, {"error": UNSUPPORTED_AUDIO_ERROR})
+    if not audio_slot.acquire(blocking=False):
+        raise HTTPException(503, "Audio processor busy", headers={"Retry-After": "5"})
+    try:
+        with tempfile.TemporaryDirectory(prefix="scamshield-upload-") as directory:
+            filename = os.path.join(directory, f"audio.{kind.extension}")
+            with open(filename, "wb") as audio:
+                audio.write(data)
+            try:
+                transcript = audio_transcript(filename)
+            except InvalidAudio as exc:
+                raise HTTPException(422, str(exc)) from exc
+            return assess(transcript)
+    finally:
+        audio_slot.release()
+
+
+@router.post(
+    "/audio", summary="Analyze an audio file for scam risk", dependencies=[Depends(authorize)],
+    responses={413: {"description": "Upload too large"}, 415: {"description": "Unsupported audio"},
+               429: {"description": "Rate limit exceeded"}, 503: {"description": "Audio processor busy"}},
+)
+def audio_check(file: UploadFile) -> riskAssessment:
+    return analyze_audio(file.file.read(MAX_FILE_SIZE + 1))
+
 
 @router.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket)-> riskAssessment | str | None:
-    """Analyze streaming audio over a WebSocket connection.
+async def websocket_endpoint(websocket: WebSocket):
+    """Analyze complete audio clips over an authenticated WebSocket connection.
 
-    The client sends audio bytes. Supported audio is transcribed and analyzed,
-    while oversized or unsupported payloads receive an error response.
+    Each binary message is one independently decodable clip, not a partial stream.
     """
+    identity = authorize_websocket(websocket)
+    origin = websocket.headers.get("origin")
+    if origin and origin not in CORS_ORIGINS:
+        await websocket.close(code=1008, reason="Origin not allowed")
+        return
     await websocket.accept()
-
     try:
         while True:
-
-            byte = await websocket.receive_bytes()
-
-            assert websocket.client is not None
-            if not rate_limit.check_rate_limit(websocket.client.host):
-                raise WebSocketException(code = 1008, reason="Reached your limit, wait 60 seconds before requesting again")
-            if len(byte) > MAX_FILE_SIZE:
-                await websocket.send_json({"error": f"File too large. Max size is {MAX_FILE_SIZE // (1024 * 1024)}MB."})
-                continue
-            kind = filetype.guess(byte) 
-
-            if kind is None or kind.mime not in Allowed:
-                await websocket.send_json({"error": UNSUPPORTED_AUDIO_ERROR})
-                continue
-            tmp_dir = "/dev/shm/" if os.path.exists("/dev/shm") else ""
-            filename = f"{tmp_dir}audio{uuid.uuid4()}.{kind.extension}"
-            with open(filename, "wb") as f:
-                f.write(byte)
-            transcript = await asyncio.to_thread(audio_transcript,filename)
-            os.remove(filename)
-            assessment =  await asyncio.to_thread(get_assessment,transcript)
-            if assessment is None: 
-                await websocket.send_json({"error": "Failed to analyze the audio. Try again"})
-            else:
-                await websocket.send_json(assessment.model_dump())
+            try:
+                data = await asyncio.wait_for(websocket.receive_bytes(), timeout=60)
+                await asyncio.to_thread(enforce_limit, identity)
+                result = await asyncio.to_thread(analyze_audio, data)
+                await websocket.send_json(result.model_dump())
+            except HTTPException as exc:
+                error = exc.detail if isinstance(exc.detail, dict) else {"error": exc.detail}
+                await websocket.send_json(error)
+                if exc.status_code in (429, 503):
+                    await websocket.close(code=1013)
+                    return
+            except asyncio.TimeoutError:
+                await websocket.close(code=1000, reason="Idle timeout")
+                return
+            except KeyError:
+                await websocket.close(code=1003, reason="Send binary audio messages")
+                return
     except WebSocketDisconnect:
-        print("Client disconnected")
+        return

app/config.py

@@ -1,26 +1,46 @@
 import os
-from openai import OpenAI
+from functools import lru_cache
+
 from dotenv import load_dotenv
+from openai import OpenAI
 
 load_dotenv()
-MAX_REQUEST_LIMIT = 10
-RATE_LIMIT_WINDOW = 60
 
-# Max allowed upload size for audio endpoints (e.g. voice notes).
-# Configurable via env var so it isn't a hardcoded magic number scattered across files.
-MAX_FILE_SIZE = int(os.environ.get(
-    "MAX_FILE_SIZE", 25 * 1024 * 1024))  # 25 MB default
+
+def positive_int(name: str, default: int) -> int:
+    value = int(os.getenv(name, str(default)))
+    if value <= 0:
+        raise ValueError(f"{name} must be positive")
+    return value
 
 
-# Allowed CORS origins
+MAX_REQUEST_LIMIT = positive_int("MAX_REQUEST_LIMIT", 10)
+RATE_LIMIT_WINDOW = positive_int("RATE_LIMIT_WINDOW", 60)
+MAX_FILE_SIZE = positive_int("MAX_FILE_SIZE", 25 * 1024 * 1024)
+MAX_AUDIO_SECONDS = positive_int("MAX_AUDIO_SECONDS", 120)
+REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")
+API_KEYS = tuple(
+    key.strip() for key in os.getenv("SCAMSHIELD_API_KEYS", "").split(",") if key.strip()
+)
 CORS_ORIGINS = [
-    origin.strip()
-    for origin in os.getenv("CORS_ORIGINS", "").split(",")
-    if origin.strip()
+    origin.strip() for origin in os.getenv("CORS_ORIGINS", "").split(",") if origin.strip()
 ]
 if any("*" in origin for origin in CORS_ORIGINS):
     raise ValueError("CORS_ORIGINS must contain explicit origins; wildcards are not allowed.")
 
-client = OpenAI(
-    api_key=os.environ.get('DEEPSEEK_API_KEY'),
-    base_url="https://api.deepseek.com")
+
+def validate_settings():
+    if not API_KEYS or any(len(key) < 32 for key in API_KEYS):
+        raise ValueError("SCAMSHIELD_API_KEYS must contain tokens of at least 32 characters")
+    if not os.getenv("DEEPSEEK_API_KEY"):
+        raise ValueError("DEEPSEEK_API_KEY is required")
+
+
+@lru_cache(maxsize=1)
+def get_ai_client() -> OpenAI:
+    return OpenAI(
+        api_key=os.environ["DEEPSEEK_API_KEY"],
+        base_url="https://api.deepseek.com",
+        timeout=20.0,
+        max_retries=0,
+    )

app/main.py

@@ -1,17 +1,36 @@
+from contextlib import asynccontextmanager
+import asyncio
+import os
+
 from fastapi import FastAPI
+from fastapi import HTTPException
 from fastapi.middleware.cors import CORSMiddleware
+from redis.exceptions import RedisError
 from app.api import analyze
-from app.config import CORS_ORIGINS
+from app.config import CORS_ORIGINS, validate_settings
+from app.middleware import BodyLimitMiddleware
+from app.services import rate_limit
+from app.services.transcription import load_model
+
+
+@asynccontextmanager
+async def lifespan(_app):
+    validate_settings()
+    if os.getenv("PRELOAD_WHISPER", "false").lower() == "true":
+        await asyncio.to_thread(load_model)
+    yield
 
 app = FastAPI(
+    lifespan=lifespan,
     title="ScamShield API",
     description=(
-        "Analyze email and audio content for scam risk. "
+        "Analyze text, email and audio content for scam risk. "
         "The API provides HTTP endpoints for email and audio analysis "
         "plus a WebSocket endpoint for streaming audio analysis."
     ),
 )
 
+app.add_middleware(BodyLimitMiddleware)
 
 app.add_middleware(
     CORSMiddleware,
@@ -27,3 +46,12 @@
 @app.get("/health")
 async def health_check():
     return {"status": "ok"}
+
+
+@app.get("/ready", include_in_schema=False)
+def readiness():
+    try:
+        rate_limit.r.ping()
+    except RedisError as exc:
+        raise HTTPException(503, "Redis unavailable") from exc
+    return {"status": "ready"}

app/middleware.py

@@ -0,0 +1,59 @@
+from fastapi import HTTPException
+from starlette.formparsers import MultiPartException
+from starlette.responses import JSONResponse
+
+from app.config import MAX_FILE_SIZE
+from app.services.auth import identify
+
+
+class BodyLimitMiddleware:  # pylint: disable=too-few-public-methods
+    """Count streamed request bytes, including requests without Content-Length."""
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+        limit = MAX_FILE_SIZE + 65536 if scope["path"] == "/audio" else 65536
+        headers = dict(scope["headers"])
+        if scope["path"] in {"/text", "/email", "/audio"}:
+            authorization = headers.get(b"authorization", b"").decode("latin-1")
+            scheme, _, token = authorization.partition(" ")
+            if scheme.lower() != "bearer" or identify(token) is None:
+                response = JSONResponse(
+                    {"detail": "Invalid or missing bearer token"}, 401,
+                    headers={"WWW-Authenticate": "Bearer"},
+                )
+                return await response(scope, receive, send)
+        try:
+            length = int(headers.get(b"content-length", b"0"))
+            if length < 0:
+                raise ValueError
+        except ValueError:
+            response = JSONResponse({"detail": "Invalid Content-Length"}, 400)
+            return await response(scope, receive, send)
+        if length > limit:
+            response = JSONResponse({"detail": "Request body too large"}, 413)
+            return await response(scope, receive, send)
+        total = 0
+
+        async def bounded_receive():
+            nonlocal total
+            message = await receive()
+            if message["type"] == "http.request":
+                total += len(message.get("body", b""))
+                if total > limit:
+                    if b"multipart/form-data" in headers.get(b"content-type", b"").lower():
+                        # Starlette closes partially spooled files for this exception type.
+                        raise MultiPartException("Request body too large")
+                    raise HTTPException(413, "Request body too large")
+            return message
+
+        async def bounded_send(message):
+            if total > limit and message["type"] == "http.response.start":
+                # Starlette maps multipart parser failures to 400; size failures are 413.
+                message = {**message, "status": 413}
+            await send(message)
+
+        return await self.app(scope, bounded_receive, bounded_send)