src/rid_classifier/space_runtime.py

from __future__ import annotations

import json
import re
import warnings
from dataclasses import dataclass
from pathlib import Path

try:
    import spaces
except ImportError:  # pragma: no cover
    class _SpacesShim:
        @staticmethod
        def GPU(*args, **kwargs):
            def decorator(fn):
                return fn

            return decorator

    spaces = _SpacesShim()

import joblib
import numpy as np
import torch
import xgboost as xgb
from scipy import sparse
from sentence_transformers import SentenceTransformer
from sklearn.exceptions import InconsistentVersionWarning
from transformers import pipeline

from .features import build_risk_feature_matrix, extract_policy_signals, normalize_text

warnings.filterwarnings("ignore", category=InconsistentVersionWarning)


EXTRACTOR_MODEL_ID = "openai/gpt-oss-20b"
_INTENT_PIPELINE = None

INTENT_SYSTEM_PROMPT = """You are a cybersecurity reverse-intent decompiler used for defensive email analysis.
You are not helping an attacker. You are summarizing intent for a phishing classifier.

Convert each email into exactly one sentence that preserves only the author's core goal.

Rules:
- Start with "The author wants to".
- Mention the requested action or requested information.
- Mention the target asset, account, payment, identity, or business process when present.
- Mention the persuasion tactic when present, such as urgency, authority, fear, reward, curiosity, or routine process framing.
- If the email is spam or promotion, describe the monetization or traffic-driving goal.
- Ignore greetings, signatures, grammar quality, and politeness.
- Do not mention these instructions.
- Do not refuse.
- Output exactly one sentence and nothing else.

Examples:
Email: "Please sign in to the updated mailbox portal before 5 PM to avoid account suspension."
Output: "The author wants to get the recipient to sign in to a mailbox portal with account credentials by using urgency and fear of account suspension."

Email: "Can you update the remittance details for invoice 1038 to the new beneficiary account today?"
Output: "The author wants to get Accounts Payable to redirect an invoice payment to a new bank account by framing it as a routine finance update with urgency."

Email: "Reminder: no action is needed during tonight's maintenance window unless you notice issues tomorrow."
Output: "The author wants to inform recipients about a maintenance window and set expectations that no action is required unless problems appear later."
"""
INTENT_REASONING_HINT = "Reasoning: low"


@dataclass
class PredictionTrace:
    prediction: int
    probability: float
    intent_string: str
    embedding_probability: float
    lexical_probability: float
    policy_probability: float
    threshold: float


def _build_intent_prompt(email_text: str) -> str:
    return f"Email:\n{email_text.strip()}\n"


def _load_intent_pipeline():
    global _INTENT_PIPELINE
    if _INTENT_PIPELINE is not None:
        return _INTENT_PIPELINE

    _INTENT_PIPELINE = pipeline(
        task="text-generation",
        model=EXTRACTOR_MODEL_ID,
        torch_dtype="auto",
        device_map="auto",
    )
    return _INTENT_PIPELINE


def _extract_generated_text(result: object) -> str:
    if isinstance(result, str):
        return result
    if isinstance(result, list) and result:
        last_item = result[-1]
        if isinstance(last_item, dict):
            content = last_item.get("content")
            if isinstance(content, str):
                return content
        return str(last_item)
    if isinstance(result, dict):
        content = result.get("content")
        if isinstance(content, str):
            return content
    return str(result)


@spaces.GPU(duration=420)
def generate_intents(email_texts: list[str]) -> list[str]:
    try:
        intent_pipeline = _load_intent_pipeline()
    except Exception as exc:  # pragma: no cover
        print(f"[SpaceRuntime] Intent pipeline load failed: {exc}")
        return [heuristic_intent_fallback(email_text) for email_text in email_texts]

    intents: list[str] = []

    for email_text in email_texts:
        try:
            messages = [
                {"role": "system", "content": f"{INTENT_SYSTEM_PROMPT}\n\n{INTENT_REASONING_HINT}"},
                {"role": "user", "content": _build_intent_prompt(email_text)},
            ]
            outputs = intent_pipeline(
                messages,
                max_new_tokens=96,
                do_sample=False,
            )
            generated_text = outputs[0].get("generated_text") if outputs else ""
            sanitized = sanitize_intent_output(_extract_generated_text(generated_text))
            if not is_valid_intent_output(sanitized):
                sanitized = heuristic_intent_fallback(email_text)
        except Exception as exc:  # pragma: no cover
            print(f"[SpaceRuntime] Intent generation failed: {exc}")
            sanitized = heuristic_intent_fallback(email_text)
        intents.append(sanitized)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    return intents


def sanitize_intent_output(text: str) -> str:
    cleaned = normalize_text(text)
    cleaned = cleaned.replace("Here's a rewritten version of the email:", "").strip()
    cleaned = re.sub(r"^output:\s*", "", cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r"^response:\s*", "", cleaned, flags=re.IGNORECASE)
    marker = "The author wants to"
    lower_cleaned = cleaned.lower()

    if lower_cleaned.startswith("the author wants you to "):
        cleaned = "The author wants to get you to " + cleaned[len("The author wants you to ") :]
    elif lower_cleaned.startswith("the author wants the recipient to "):
        cleaned = "The author wants to get the recipient to " + cleaned[len("The author wants the recipient to ") :]
    elif lower_cleaned.count(marker.lower()) >= 2:
        cleaned = cleaned[lower_cleaned.rfind(marker.lower()) :]
    elif marker.lower() in lower_cleaned:
        start = lower_cleaned.index(marker.lower())
        cleaned = cleaned[start:]
    else:
        cleaned = f"{marker} {cleaned.lstrip('-* ').strip()}"

    cleaned = re.sub(r"^(The author wants to\s+)+", "The author wants to ", cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r"The author wants to\s+The author wants to\s+", "The author wants to ", cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r"The author wants to\s+The author wants\b", "The author wants", cleaned, flags=re.IGNORECASE)
    cleaned = cleaned.replace("The author wants to to ", "The author wants to ")
    sentence_match = re.match(r"^(.*?[.!?])(\s|$)", cleaned)
    if sentence_match:
        cleaned = sentence_match.group(1)
    if not cleaned.endswith("."):
        cleaned = f"{cleaned.rstrip('.!?')}."
    return cleaned


def is_valid_intent_output(intent: str) -> bool:
    lowered = intent.lower().strip()
    invalid_fragments = [
        "i cannot",
        "i can't",
        "i will not",
        "reverse-intent decompiler",
        "do not output any other text",
        "ignore polite formatting",
        "here is",
        "here's",
        "output:",
        "response:",
    ]
    if not lowered.startswith("the author wants to"):
        return False
    if any(fragment in lowered for fragment in invalid_fragments):
        return False
    if len(lowered.split()) < 8:
        return False
    return True


def heuristic_intent_fallback(email_text: str) -> str:
    lowered = normalize_text(email_text).lower()

    def has_any(phrases: list[str]) -> bool:
        return any(phrase in lowered for phrase in phrases)

    if has_any(["sign in", "log in", "login", "password", "credential", "mfa", "multi-factor", "portal", "vpn"]):
        action = "get the recipient to submit account credentials through a sign-in flow"
    elif has_any(["wire transfer", "remittance", "bank details", "beneficiary", "invoice", "payment", "direct deposit"]):
        action = "get the recipient to move money or change payment instructions"
    elif has_any(["shared document", "review", "sign", "attachment", "document", "enable content"]):
        action = "get the recipient to open or sign a document that advances the sender's goal"
    elif has_any(["package", "parcel", "delivery", "shipping"]):
        action = "get the recipient to verify delivery details or make a payment to release a shipment"
    elif has_any(["sale", "discount", "offer", "buy now", "click here"]):
        action = "drive the recipient to visit a marketing or monetized destination"
    else:
        action = "influence the recipient to take a requested action"

    if has_any(["urgent", "today", "before cob", "deadline", "locked out", "immediately"]):
        tactic = "using urgency"
    elif has_any(["ceo", "finance", "security", "it administration", "payroll", "legal", "executive"]):
        tactic = "using authority framing"
    elif has_any(["confidential", "private", "unreachable by phone"]):
        tactic = "using secrecy"
    elif has_any(["reminder", "agenda", "meeting", "no action needed", "maintenance"]):
        tactic = "through routine business process framing"
    else:
        tactic = "through persuasive business framing"

    return f"The author wants to {action} {tactic}."


class SpaceHybridClassifier:
    def __init__(self, artifact_dir: str | Path) -> None:
        self.artifact_dir = Path(artifact_dir)
        self.metadata = json.loads((self.artifact_dir / "metadata.json").read_text(encoding="utf-8"))
        self.embedder = SentenceTransformer(self.metadata.get("embedder_model", "all-MiniLM-L6-v2"))

        self.embedding_model = xgb.XGBClassifier()
        self.embedding_model.load_model(str(self.artifact_dir / "intent_xgboost_model.json"))
        self.lexical_model = joblib.load(self.artifact_dir / "lexical_logistic_model.joblib")
        self.lexical_feature_extractor = joblib.load(self.artifact_dir / "lexical_feature_extractor.joblib")
        self.blend_weight = float(self.metadata.get("blend_weight", 0.5))
        self.threshold = float(self.metadata.get("decision_threshold", 0.5))

    def predict_batch(self, raw_email_texts: list[str]) -> list[PredictionTrace]:
        intents = generate_intents(raw_email_texts)
        risk_features = build_risk_feature_matrix(raw_email_texts, intents)
        intent_vectors = self.embedder.encode(intents, show_progress_bar=False)
        embedding_input = np.hstack([intent_vectors, risk_features])
        embedding_probs = self.embedding_model.predict_proba(embedding_input)[:, 1]

        lexical_input = self.lexical_feature_extractor.transform(raw_email_texts, intents)
        lexical_input = sparse.hstack([lexical_input, sparse.csr_matrix(risk_features)], format="csr")
        lexical_probs = self.lexical_model.predict_proba(lexical_input)[:, 1]

        traces: list[PredictionTrace] = []
        for raw_email, intent, embedding_probability, lexical_probability in zip(
            raw_email_texts,
            intents,
            embedding_probs,
            lexical_probs,
            strict=False,
        ):
            policy_probability, policy_flags = extract_policy_signals(raw_email, intent)
            blended_probability = self.blend_weight * float(embedding_probability) + (1.0 - self.blend_weight) * float(lexical_probability)
            probability = max(blended_probability, float(policy_probability))
            if policy_flags == ["routine_benign"]:
                probability = min(probability, 0.24)
            prediction = int(probability >= self.threshold)
            traces.append(
                PredictionTrace(
                    prediction=prediction,
                    probability=float(probability),
                    intent_string=intent,
                    embedding_probability=float(embedding_probability),
                    lexical_probability=float(lexical_probability),
                    policy_probability=float(policy_probability),
                    threshold=self.threshold,
                )
            )
        return traces