phishing.py

import json
import os
import random
from typing import Any, Dict, Optional

import httpx
from fastapi import Header, HTTPException, Request
from fastapi.responses import JSONResponse
from helper.ratelimit import enforce_rate_limit
from helper.subscriptions import resolve_token_identity

from . import router


SHIELD_API_URL = "https://api.cerebras.ai/v1/chat/completions"
SHIELD_ANALYSIS_MODEL = "gpt-oss-120b"


def _get_shield_api_key() -> Optional[str]:
    raw = os.getenv("CER_KEY", "")
    keys = [k.strip() for k in raw.split(",") if k.strip()]
    return random.choice(keys) if keys else None


PHISHING_SYSTEM_PROMPT = """You are a phishing detection analyst.
Analyze the provided content for phishing indicators.
Return JSON with:
- is_phishing (boolean)
- confidence (0-1)
- risk_score (0-100)
- indicators (list of observed phishing indicators)
- threat_level (low, medium, high, critical)
- explanation (brief reason for the verdict)"""


def _parse_phishing_response(raw: str) -> Optional[Dict[str, Any]]:
    try:
        data = json.loads(raw)
    except json.JSONDecodeError:
        try:
            start = raw.index("{")
            end = raw.rindex("}") + 1
            data = json.loads(raw[start:end])
        except (ValueError, json.JSONDecodeError):
            return None

    if not isinstance(data, dict):
        return None

    return {
        "is_phishing": bool(data.get("is_phishing", False)),
        "confidence": max(0.0, min(1.0, float(data.get("confidence", 0)))),
        "risk_score": max(0, min(100, int(data.get("risk_score", 0)))),
        "indicators": data.get("indicators", []) if isinstance(data.get("indicators"), list) else [],
        "threat_level": data.get("threat_level", "low") if data.get("threat_level") in ("low", "medium", "high", "critical") else "low",
        "explanation": data.get("explanation", ""),
    }


@router.post("/detect-phishing")
async def detect_phishing(
    request: Request,
    authorization: Optional[str] = Header(None),
    x_client_id: Optional[str] = Header(None),
) -> JSONResponse:
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")

    token = authorization.split(" ", 1)[1].strip()
    identity = await resolve_token_identity(token)
    if not identity:
        raise HTTPException(status_code=401, detail="Invalid authorization token")

    await enforce_rate_limit(request, authorization, "aiShieldDaily", x_client_id)

    try:
        body = await request.json()
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid JSON body")

    if not isinstance(body, dict):
        raise HTTPException(status_code=400, detail="Request body must be a JSON object")

    content = body.get("content", "")
    url = body.get("url", "")
    sender = body.get("sender", "")
    subject = body.get("subject", "")

    if not content and not url and not sender and not subject:
        raise HTTPException(status_code=400, detail="Provide at least one of: content, url, sender, subject")

    analysis_input = {}
    if content:
        analysis_input["content"] = content
    if url:
        analysis_input["url"] = url
    if sender:
        analysis_input["sender"] = sender
    if subject:
        analysis_input["subject"] = subject

    user_content = json.dumps(analysis_input, indent=2)

    payload = {
        "model": SHIELD_ANALYSIS_MODEL,
        "messages": [
            {"role": "system", "content": PHISHING_SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
        ],
        "temperature": 0.1,
        "max_tokens": 1024,
        "response_format": {"type": "json_object"},
    }

    api_key = _get_shield_api_key()
    if not api_key:
        raise HTTPException(status_code=502, detail="No API key configured")

    try:
        async with httpx.AsyncClient(timeout=30) as client:
            r = await client.post(
                SHIELD_API_URL,
                json=payload,
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
            )
            if r.status_code >= 400:
                raise HTTPException(status_code=502, detail="Upstream model error")

            body_resp = r.json()
            choice = body_resp.get("choices", [{}])[0]
            raw = choice.get("message", {}).get("content", "")
            result = _parse_phishing_response(raw)
            if not result:
                raise HTTPException(status_code=502, detail="Failed to parse model response")

            return JSONResponse(result)
    except HTTPException:
        raise
    except Exception:
        raise HTTPException(status_code=502, detail="Upstream model error")