main.py

"""
DeepGuard — FastAPI Backend
Real-time, stateless deepfake detection API.

Endpoints:
  GET  /health   — Liveness check (used by extension popup)
  POST /analyze  — Analyze an image for AI-generation artifacts

All data is processed in RAM and dropped immediately after the response.
"""

import io
import base64
import traceback
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from PIL import Image

import inference
import metadata as meta_module
import heatmap as heatmap_module


# ---------------------------------------------------------------------------
# Lifespan: load model once at startup
# ---------------------------------------------------------------------------

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("[DeepGuard] Starting up — loading ONNX model...")
    try:
        inference.load_model()
        print("[DeepGuard] Model ready. Server is live at http://localhost:8000")
    except FileNotFoundError as e:
        print(f"[DeepGuard] WARNING: {e}")
        print("[DeepGuard] Run 'python download_model.py' to fetch the model.")
    yield
    print("[DeepGuard] Shutting down.")


# ---------------------------------------------------------------------------
# App setup
# ---------------------------------------------------------------------------

app = FastAPI(
    title="DeepGuard API",
    description="Real-time stateless deepfake detection using ViT + ONNX Runtime",
    version="1.0.0",
    lifespan=lifespan,
)

# Aggressive CORS — required because Chrome extensions use a chrome-extension:// origin
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


import ela as ela_module


# ---------------------------------------------------------------------------
# Request / Response schemas
# ---------------------------------------------------------------------------

class AnalyzeRequest(BaseModel):
    image_data: str  # data:image/...;base64,<payload>  OR  raw base64


class MetadataResult(BaseModel):
    exif_data_present: bool
    software_signature_found: str
    warning: str


class ForensicsResult(BaseModel):
    model_reasoning: str
    metadata: MetadataResult


class AnalyzeResponse(BaseModel):
    status: str
    threat_level: str
    confidence_score: float
    heatmap_overlay_url: str
    ela_overlay_url: str
    forensics: ForensicsResult


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _decode_image(image_data: str) -> tuple[bytes, Image.Image]:
    """
    Decode a base64 data-URI or raw base64 string into (raw_bytes, PIL Image).
    """
    if image_data.startswith("data:"):
        # Strip "data:image/jpeg;base64," prefix
        header, b64_str = image_data.split(",", 1)
    else:
        b64_str = image_data

    try:
        raw_bytes = base64.b64decode(b64_str)
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid base64 image data.")

    try:
        image = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="Could not decode image from base64 payload.")

    return raw_bytes, image


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------

@app.get("/health")
async def health():
    """Liveness check. Returns model load status."""
    try:
        session = inference.get_session()
        return {
            "status": "ok",
            "model_loaded": True,
            "attention_heatmap": inference.has_attention_outputs(),
        }
    except RuntimeError:
        return {
            "status": "degraded",
            "model_loaded": False,
            "attention_heatmap": False,
            "message": "Model not loaded. Run python download_model.py",
        }


@app.post("/analyze", response_model=AnalyzeResponse)
async def analyze(request: AnalyzeRequest):
    """
    Main analysis endpoint.
    Accepts a base64-encoded image and returns:
      - Deepfake confidence score
      - Threat level classification
      - Grad-CAM / attention heatmap overlay (base64 PNG)
      - EXIF metadata forensics
    """
    # ── 1. Decode image ──────────────────────────────────────────────────
    try:
        raw_bytes, image = _decode_image(request.image_data)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Image decode error: {e}")

    # ── 2. Metadata forensics (DFIR) ─────────────────────────────────────
    try:
        forensic_meta = meta_module.extract_metadata(raw_bytes)
    except Exception:
        forensic_meta = {
            "exif_data_present": False,
            "software_signature_found": "None",
            "warning": "Metadata extraction failed.",
        }

    # ── 3. AI inference ───────────────────────────────────────────────────
    try:
        confidence_score, output_dict = inference.run_inference(image)
    except RuntimeError as e:
        raise HTTPException(
            status_code=503,
            detail=f"Model not loaded: {e}. Run python download_model.py first.",
        )
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Inference error: {e}")

    # ── 4. Heatmap & ELA generation ────────────────────────────────────────
    try:
        heatmap_url = heatmap_module.generate_heatmap(image, output_dict, confidence_score)
    except Exception:
        import traceback
        traceback.print_exc()
        # Fallback: return a 1×1 transparent PNG
        heatmap_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="

    # Run Error Level Analysis (ELA)
    ela_url = ela_module.generate_ela(image)

    # ── 5. Build response ─────────────────────────────────────────────────
    threat_level = inference.get_threat_level(confidence_score)
    model_reasoning = inference.get_model_reasoning(
        confidence_score,
        forensic_meta["exif_data_present"],
        forensic_meta["software_signature_found"],
    )

    return AnalyzeResponse(
        status="success",
        threat_level=threat_level,
        confidence_score=round(confidence_score, 4),
        heatmap_overlay_url=heatmap_url,
        ela_overlay_url=ela_url,
        forensics=ForensicsResult(
            model_reasoning=model_reasoning,
            metadata=MetadataResult(**forensic_meta),
        ),
    )