Version_3/main.py

import math
import os
from contextlib import asynccontextmanager

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

# ── Import your existing cli.py functions directly ───────────────────────────
# If running from inside Version_3/ folder this just works.
# If your root main.py imports this, adjust the import to:
#   from Server_Main.Version_3.cli import clean_text, predict, load_models
from cli import clean_text, predict, load_models

# ─────────────────────────────────────────────────────────────────────────────
# STARTUP: load models once when the server boots, not on every request
# ─────────────────────────────────────────────────────────────────────────────
_vectorizer = None
_model      = None
_load_error = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load the SVM model and vectorizer at startup."""
    global _vectorizer, _model, _load_error
    try:
        _vectorizer, _model = load_models()
        print("[Chimera] Models loaded OK.")
    except SystemExit:
        # load_models() calls sys.exit(1) if files are missing —
        # we catch that here so the server still starts and returns
        # a clean error response instead of crashing.
        _load_error = "Model files not found. Check models/SVM_model.pkl and models/vectorizer.pkl exist."
        print(f"[Chimera] WARNING: {_load_error}")
    except Exception as e:
        _load_error = str(e)
        print(f"[Chimera] WARNING: Could not load models — {_load_error}")
    yield  # server runs here
    # (nothing to clean up on shutdown)


# ─────────────────────────────────────────────────────────────────────────────
# APP SETUP
# ─────────────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Chimera Email Scanner",
    description="SVM-based phishing / spam email detection API",
    version="3.0.0",
    lifespan=lifespan,
)

# Allow your frontend origin — update the list if you host the HTML elsewhere
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],   # tighten this to your domain in production
    allow_methods=["POST", "GET"],
    allow_headers=["*"],
)


# ─────────────────────────────────────────────────────────────────────────────
# REQUEST / RESPONSE MODELS
# ─────────────────────────────────────────────────────────────────────────────
class EmailScanRequest(BaseModel):
    """
    Matches the payload chimera-api.js sends in checkEmail():

        {
          message_id:    "msg_1234567890",
          sender_domain: "badguy.com",
          display_name:  "Subject line goes here",
          raw_text:      "Full email body text..."
        }
    """
    message_id:    str = "msg_unknown"
    sender_domain: str = "unknown.com"
    display_name:  str = "Unknown Sender"
    raw_text:      str


class EmailScanResponse(BaseModel):
    risk_score:  float         # 0 – 100  (higher = more dangerous)
    verdict:     str           # human-readable label for the badge
    label:       str           # raw model output: "Spam" or "Ham"
    confidence:  float | None  # predict_proba spam probability %, or null
    details:     dict          # three string fields rendered in the grid cards


# ─────────────────────────────────────────────────────────────────────────────
# HELPERS
# ─────────────────────────────────────────────────────────────────────────────
def _verdict_label(spam_score: float) -> str:
    """Map spam_score (0-100) to the verdict string shown in the frontend badge."""
    if spam_score >= 70:
        return "PHISHING"
    if spam_score >= 40:
        return "SUSPICIOUS"
    return "SAFE"


# ─────────────────────────────────────────────────────────────────────────────
# ROUTES
# ─────────────────────────────────────────────────────────────────────────────
@app.get("/")
def health_check():
    """Simple liveness probe — HuggingFace hits this to confirm the Space is up."""
    status = "ok" if _model is not None else "degraded"
    return {"status": status, "service": "Chimera Email Scanner v3"}


@app.post("/email", response_model=EmailScanResponse)
def scan_email(req: EmailScanRequest):
    """
    Analyse an email for spam / phishing using the SVM model from cli.py.
    """
    global _model, _vectorizer, _load_error

    # --- FIX: Lazy-Loading for Mounted Apps ---
    # Because FastAPI ignores lifespan events on mounted apps,
    # we load the models on the first request if they are missing.
    if _model is None or _vectorizer is None:
        try:
            print("[Chimera] Lazy-loading SVM models on first request...")
            _vectorizer, _model = load_models()
        except SystemExit:
            _load_error = "Model files not found in Version_3/models/"
        except Exception as e:
            _load_error = str(e)

    # Guard: models didn't load
    if _model is None or _vectorizer is None:
        error_msg = _load_error or "Models not loaded."
        return EmailScanResponse(
            risk_score  = 0,
            verdict     = "Model Unavailable",
            label       = "Unknown",
            confidence  = None,
            details     = {
                "language_analysis": f"Server error: {error_msg}",
                "link_analysis":     "Cannot analyse — model unavailable.",
                "sender_check":      f"Domain received: {req.sender_domain}",
            }
        )

    # Guard: empty body
    if not req.raw_text.strip():
        raise HTTPException(status_code=422, detail="raw_text must not be empty.")

    # ── Run prediction (reuses cli.py's predict() directly) ──────────────────
    combined_text = f"{req.display_name}\n{req.raw_text}"

    label, confidence, spam_score = predict(combined_text, _vectorizer, _model)

    risk_score = spam_score if spam_score is not None else (100.0 if label == "Spam" else 0.0)
    risk_score = round(risk_score, 1)

    # ── Build the three detail strings the frontend grid cards display ────────
    conf_str = f"{confidence:.1f}%" if confidence is not None else "N/A"

    if label == "Spam":
        language_msg = (
            f"Spam language detected — risk score {risk_score:.0f}% "
            f"(model confidence: {conf_str})"
        )
        link_msg = (
            "Treat any embedded links with caution — "
            "spam/phishing content identified in body"
        )
        sender_msg = (
            f"Sender domain '{req.sender_domain}' flagged — "
            "matches spam characteristics"
        )
    else:
        language_msg = (
            f"No spam patterns found — risk score {risk_score:.0f}% "
            f"(model confidence: {conf_str})"
        )
        link_msg     = "No suspicious link patterns detected in email content"
        sender_msg   = (
            f"Sender domain '{req.sender_domain}' — "
            "no issues detected"
        )

    return EmailScanResponse(
        risk_score = risk_score,
        verdict    = _verdict_label(risk_score),
        label      = label,
        confidence = confidence,
        details    = {
            "language_analysis": language_msg,
            "link_analysis":     link_msg,
            "sender_check":      sender_msg,
        }
    )