analyze.py

import asyncio
import json
import os
import random
from typing import Any, Dict, List, 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
from .agents import (
    analyze_content_threats,
    analyze_exfiltration,
    analyze_fraud,
    analyze_identity,
    analyze_prompt_threats,
    plan_investigation,
)
from .decision import build_decision
from .evidence import aggregate_evidence
from .graph import GraphStore, detect_campaign, detect_duplicates
from .heuristics import run_heuristics
from .intelligence import (
    collect_behavior_evidence,
    collect_content_evidence,
    collect_device_evidence,
    collect_email_evidence,
    collect_ip_evidence,
    collect_phone_evidence,
    collect_prompt_evidence,
    collect_username_evidence,
)
from .memory import CampaignStore, CustomerMemory, GlobalMemory
from .models.evidence import Evidence
from .utils import entity_values, hash_entity, normalize_signals

DEFAULT_CONFIG = {
    "features": {
        "heuristics": True,
        "historical_memory": True,
        "duplicate_detection": False,
        "campaign_detection": False,
        "email_intelligence": True,
        "ip_intelligence": True,
        "phone_intelligence": True,
        "username_intelligence": True,
        "device_intelligence": True,
        "behavior_intelligence": True,
        "content_intelligence": True,
        "prompt_intelligence": True,
        "identity_analysis": True,
        "fraud_analysis": True,
        "prompt_analysis": True,
        "content_analysis": True,
        "exfiltration_analysis": True,
        "llm_reasoning": True,
        "memory_update": False,
    }
}


def resolve_config(body: Dict[str, Any]) -> Dict[str, Any]:
    given = body.get("config", {})
    features = dict(DEFAULT_CONFIG["features"])
    given_features = given.get("features", {}) if isinstance(given, dict) else {}
    if isinstance(given_features, dict):
        for key, value in given_features.items():
            if key in features and isinstance(value, bool):
                features[key] = value
    return {"features": features}

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

SYSTEM_PROMPT = """You are a senior fraud analyst.
You must only reason using supplied evidence. Do not invent evidence.
Return JSON with:
- risk_score (0-100)
- confidence (0-1)
- decision (allow, challenge, rate_limit, review, block)
- threat_categories (list)
- reasons (list of short explanations)
- recommended_action
Use relationship, historical, campaign, and specialized threat evidence when present."""


def _build_llm_payload(
    signals: Dict[str, Any],
    evidence: List[Evidence],
    duplicate_analysis: Dict[str, Any],
    historical_intelligence: Dict[str, Any],
    campaign_analysis: Dict[str, Any],
) -> Dict[str, Any]:
    signal_summary = {
        "has_email": bool(signals.get("email")),
        "has_phone": bool(signals.get("phone")),
        "has_username": bool(signals.get("username")),
        "has_ip": bool(signals.get("ip")),
        "has_device": bool(signals.get("device_fingerprint")),
        "has_content": bool(signals.get("content")),
        "content_length": len(signals.get("content") or ""),
        "has_metadata": isinstance(signals.get("metadata"), dict),
    }
    campaign_summary = {
        "campaign_id": campaign_analysis.get("campaign_id"),
        "campaign_risk_score": campaign_analysis.get("campaign_risk_score", 0),
        "confidence": campaign_analysis.get("confidence", 0.0),
        "reasons": campaign_analysis.get("reasons", []),
    }
    user_content = json.dumps({
        "signals": signal_summary,
        "evidence": [item.to_dict() for item in evidence],
        "duplicate_analysis": {
            "duplicate_score": duplicate_analysis.get("duplicate_score", 0),
            "confidence": duplicate_analysis.get("confidence", 0.0),
            "linked_account_count": len(duplicate_analysis.get("linked_accounts", [])),
            "reasons": duplicate_analysis.get("reasons", []),
        },
        "historical_intelligence": historical_intelligence,
        "campaign_analysis": campaign_summary,
    }, indent=2)

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


def _parse_llm_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

    risk_score = data.get("risk_score")
    confidence = data.get("confidence")

    if risk_score is None:
        return None

    return {
        "risk_score": max(0, min(100, int(risk_score))),
        "confidence": max(0.0, min(1.0, float(confidence))) if confidence is not None else 0.5,
        "decision": data.get("decision", "review")
        if data.get("decision") in ("allow", "challenge", "rate_limit", "review", "block")
        else "review",
        "reasons": data.get("reasons", []) if isinstance(data.get("reasons"), list) else [],
        "recommended_action": data.get("recommended_action", data.get("decision", "review")),
        "threat_categories": data.get("threat_categories", [])
        if isinstance(data.get("threat_categories"), list)
        else [],
    }


async def _call_llm(payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    api_key = _get_shield_api_key()
    if not SHIELD_API_URL or not api_key:
        return None

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    try:
        async with httpx.AsyncClient(timeout=30) as client:
            r = await client.post(SHIELD_API_URL, json=payload, headers=headers)
            if r.status_code >= 400:
                return None
            body = r.json()
            choice = body.get("choices", [{}])[0]
            raw = choice.get("message", {}).get("content", "")
            return _parse_llm_response(raw)
    except Exception:
        return None


def _metadata_account_id(body: Dict[str, Any]) -> Optional[str]:
    metadata = body.get("metadata")
    if isinstance(metadata, dict):
        for key in ("account_id", "user_id", "external_user_id"):
            value = metadata.get(key)
            if isinstance(value, str) and value.strip():
                return value.strip()
    for key in ("account_id", "subject_user_id", "external_user_id"):
        value = body.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip()
    return None


def _subject_account_hash(body: Dict[str, Any], signals: Dict[str, Any]) -> str:
    account_id = _metadata_account_id(body)
    if account_id:
        return hash_entity("account", account_id)
    for key in ("email", "phone", "username", "device_fingerprint"):
        value = signals.get(key)
        if isinstance(value, str) and value:
            return hash_entity("account", f"{key}:{value}")
    return hash_entity("account", json.dumps({k: bool(v) for k, v in signals.items()}, sort_keys=True))


def _heuristic_evidence(score: int, reasons: List[str]) -> List[Evidence]:
    if not score:
        return []
    return [
        Evidence(
            source="heuristics",
            category="heuristic_risk",
            risk_score=score,
            weight=0.55,
            confidence=0.7,
            explanation=reason,
        )
        for reason in reasons
    ]


async def _safe_collect(awaitable: Any) -> List[Evidence]:
    try:
        result = await awaitable
        return result if isinstance(result, list) else []
    except Exception:
        return []


def _memory_evidence(historical: Dict[str, Any]) -> List[Evidence]:
    evidence: List[Evidence] = []
    for entity_type, record in historical.items():
        if not isinstance(record, dict):
            continue
        abuse_count = int(record.get("abuse_count", 0))
        avg_risk = float(record.get("avg_risk", 0))
        linked_customers = int(record.get("linked_customers", 0))
        if abuse_count or avg_risk >= 60:
            evidence.append(Evidence(
                source="global_memory",
                category="historical_abuse",
                risk_score=min(100, max(int(avg_risk), 35 + (abuse_count * 8))),
                weight=0.72,
                confidence=0.82,
                explanation=f"{entity_type.title()} has prior abuse history in Shield memory",
                metadata={"abuse_count": abuse_count, "linked_customers": linked_customers},
            ))
    return evidence


def _update_memory(
    *,
    customer_id: Optional[str],
    account_hash: str,
    entities: Dict[str, str],
    graph: GraphStore,
    global_memory: GlobalMemory,
    customer_memory: CustomerMemory,
    campaign_store: CampaignStore,
    campaign_analysis: Dict[str, Any],
    result: Dict[str, Any],
) -> None:
    try:
        graph.add_account_entities(customer_id=customer_id, account_hash=account_hash, entities=entities)
        global_memory.update(
            entities,
            customer_id=customer_id,
            risk_score=int(result["risk_score"]),
            decision=str(result["decision"]),
        )
        customer_memory.update(
            customer_id,
            entities,
            account_hash=account_hash,
            risk_score=int(result["risk_score"]),
            decision=str(result["decision"]),
            reasons=list(result.get("reasons", [])),
        )
        campaign_id = campaign_analysis.get("campaign_id")
        if campaign_id:
            entity_hashes = {entity_type: hash_entity(entity_type, value) for entity_type, value in entities.items()}
            campaign_store.update(
                str(campaign_id),
                account_hash=account_hash,
                entity_hashes=entity_hashes,
                risk_score=int(result["risk_score"]),
            )
    except Exception:
        return


@router.post("/analyze")
async def analyze(
    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")

    email = body.get("email")
    phone = body.get("phone")
    username = body.get("username")
    ip = body.get("ip")
    device_fingerprint = body.get("device_fingerprint")
    geolocation = body.get("geolocation")
    signup_time = body.get("signup_time")
    content = body.get("content")
    metadata = body.get("metadata")

    cfg = resolve_config(body)
    features = cfg["features"]

    signals = normalize_signals({
        "email": email,
        "phone": phone,
        "username": username,
        "ip": ip,
        "device_fingerprint": device_fingerprint,
        "geolocation": geolocation,
        "signup_time": signup_time,
        "content": content,
        "metadata": metadata if isinstance(metadata, dict) else None,
    })

    customer_id = str(identity.get("user_id") or identity.get("email") or "anonymous")
    account_hash = _subject_account_hash(body, signals)
    entities = entity_values(signals)

    plan = plan_investigation(signals)
    tools = set(plan["tools"])

    # Apply config to filter tools
    if not features["heuristics"]:
        tools.discard("heuristics")
    if not features["historical_memory"]:
        tools.discard("historical_memory")
    if not features["duplicate_detection"]:
        tools.discard("duplicate_detection")
    if not features["campaign_detection"]:
        tools.discard("campaign_detection")
    if not features["email_intelligence"]:
        tools.discard("email_intelligence")
    if not features["ip_intelligence"]:
        tools.discard("ip_intelligence")
    if not features["phone_intelligence"]:
        tools.discard("phone_intelligence")
    if not features["username_intelligence"]:
        tools.discard("username_intelligence")
    if not features["device_intelligence"]:
        tools.discard("device_intelligence")
    if not features["behavior_intelligence"]:
        tools.discard("behavior_intelligence")
    if not features["content_intelligence"]:
        tools.discard("content_intelligence")
    if not features["prompt_intelligence"]:
        tools.discard("prompt_intelligence")

    graph = GraphStore()
    global_memory = GlobalMemory()
    customer_memory = CustomerMemory()
    campaign_store = CampaignStore()

    historical_intelligence = {}
    if "historical_memory" in tools:
        historical_intelligence = global_memory.lookup(entities)

    heuristic_score = 0
    heuristic_reasons = []
    if "heuristics" in tools:
        heuristic_score, heuristic_reasons = run_heuristics(
            email=signals.get("email"),
            ip=signals.get("ip"),
            content=signals.get("content"),
        )

    evidence: List[Evidence] = _heuristic_evidence(heuristic_score, heuristic_reasons)
    evidence.extend(_memory_evidence(historical_intelligence))

    duplicate_analysis = {"duplicate_score": 0, "linked_accounts": [], "evidence": [], "reasons": []}
    if "duplicate_detection" in tools:
        duplicate_analysis = detect_duplicates(entities, account_hash=account_hash, graph=graph)
        evidence.extend(duplicate_analysis.get("evidence", []))

    campaign_analysis = {"campaign_id": None, "campaign_risk_score": 0, "evidence": [], "reasons": [], "confidence": 0.0}
    if "campaign_detection" in tools:
        campaign_analysis = detect_campaign(
            signals,
            entities,
            account_hash=account_hash,
            store=campaign_store,
        )
        evidence.extend(campaign_analysis.get("evidence", []))

    collectors = []
    if "email_intelligence" in tools:
        collectors.append(_safe_collect(collect_email_evidence(signals.get("email"))))
    if "ip_intelligence" in tools:
        collectors.append(_safe_collect(collect_ip_evidence(signals.get("ip"), historical_intelligence.get("ip"))))

    if collectors:
        for collected in await asyncio.gather(*collectors):
            evidence.extend(collected)

    if "phone_intelligence" in tools:
        evidence.extend(collect_phone_evidence(signals.get("phone")))
    if "username_intelligence" in tools:
        evidence.extend(collect_username_evidence(signals.get("username")))
    if "device_intelligence" in tools:
        evidence.extend(collect_device_evidence(signals.get("device_fingerprint"), signals.get("metadata")))
    if "behavior_intelligence" in tools:
        evidence.extend(collect_behavior_evidence(signals.get("metadata")))
    if "content_intelligence" in tools:
        evidence.extend(collect_content_evidence(signals.get("content")))
    if "prompt_intelligence" in tools:
        evidence.extend(collect_prompt_evidence(signals.get("content")))

    if features["identity_analysis"]:
        evidence.extend(analyze_identity(duplicate_analysis))
    if features["fraud_analysis"]:
        evidence.extend(analyze_fraud(evidence, duplicate_analysis, campaign_analysis))
    if features["prompt_analysis"]:
        evidence.extend(analyze_prompt_threats(evidence))
    if features["content_analysis"]:
        evidence.extend(analyze_content_threats(evidence))
    if features["exfiltration_analysis"]:
        evidence.extend(analyze_exfiltration(evidence))

    aggregate = aggregate_evidence(evidence)
    llm_result = None
    if features["llm_reasoning"] and evidence:
        llm_payload = _build_llm_payload(
            signals,
            evidence,
            duplicate_analysis,
            historical_intelligence,
            campaign_analysis,
        )
        llm_result = await _call_llm(llm_payload)

    result = build_decision(aggregate, llm_result)
    result["duplicate_user_score"] = duplicate_analysis.get("duplicate_score", 0)
    result["linked_accounts"] = len(duplicate_analysis.get("linked_accounts", []))
    result["campaign_risk_score"] = campaign_analysis.get("campaign_risk_score", 0)
    result["threat_categories"] = aggregate.get("threat_categories", [])
    if llm_result:
        for category in llm_result.get("threat_categories", []):
            if isinstance(category, str) and category not in result["threat_categories"]:
                result["threat_categories"].append(category)
    result["investigation"] = {
        "tools": sorted(tools),
        "evidence_count": len(evidence),
        "campaign_id": campaign_analysis.get("campaign_id"),
        "historical_matches": list(historical_intelligence.keys()),
    }
    result["evidence"] = [item.to_dict() for item in evidence]
    result["config_applied"] = features

    if features["memory_update"]:
        _update_memory(
            customer_id=customer_id,
            account_hash=account_hash,
            entities=entities,
            graph=graph,
            global_memory=global_memory,
            customer_memory=customer_memory,
            campaign_store=campaign_store,
            campaign_analysis=campaign_analysis,
            result=result,
        )

    return JSONResponse(result)