src/rid_classifier/features.py

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Iterable

import numpy as np
from scipy import sparse
from sklearn.feature_extraction.text import TfidfVectorizer


WHITESPACE_RE = re.compile(r"\s+")
URL_RE = re.compile(r"https?://\S+|www\.\S+", re.IGNORECASE)
EMAIL_RE = re.compile(r"\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b", re.IGNORECASE)

CREDENTIAL_PATTERNS = [
    r"\bsign in\b",
    r"\blog ?in\b",
    r"\bpassword\b",
    r"\bcredential(?:s)?\b",
    r"\bverify (?:your )?account\b",
    r"\bvalidate (?:your )?(?:vpn|profile|account)\b",
    r"\bmfa\b",
    r"\bmulti[- ]factor\b",
    r"\bre-enroll\b",
    r"\bmailbox\b",
    r"\bportal\b",
]
PAYMENT_PATTERNS = [
    r"\bwire transfer\b",
    r"\bremittance\b",
    r"\bbank (?:details|account)\b",
    r"\bbeneficiary\b",
    r"\bpayment\b",
    r"\binvoice\b",
    r"\bdirect deposit\b",
    r"\bpayroll\b",
]
URGENCY_PATTERNS = [
    r"\burgent\b",
    r"\basap\b",
    r"\bimmediately\b",
    r"\bbefore (?:cob|close of business|end of day|today)\b",
    r"\btoday\b",
    r"\bdeadline\b",
    r"\bexpire[sd]?\b",
    r"\blocked out\b",
    r"\baction required\b",
    r"\bmust\b",
]
AUTHORITY_PATTERNS = [
    r"\bit administration\b",
    r"\bsecurity operations\b",
    r"\bnetwork engineering\b",
    r"\bpayroll services\b",
    r"\bhr operations\b",
    r"\baccounts payable\b",
    r"\bexecutive assistant\b",
    r"\bexternal counsel\b",
    r"\bceo\b",
    r"\bfinance\b",
]
DOCUMENT_PATTERNS = [
    r"\battachment\b",
    r"\bshared document\b",
    r"\bopen the (?:document|file|attachment)\b",
    r"\breview\b",
    r"\bsign\b",
    r"\benable content\b",
    r"\bpdf\b",
    r"\bdoc(?:ument)?\b",
]
BENIGN_PATTERNS = [
    r"\bno action (?:is )?needed\b",
    r"\bfor awareness\b",
    r"\breminder\b",
    r"\bagenda\b",
    r"\bmeeting\b",
    r"\bworkshop\b",
    r"\bdietary\b",
    r"\bopen enrollment\b",
    r"\bbenefits\b",
    r"\bproject status\b",
    r"\bstatus document\b",
    r"\bcode freeze\b",
    r"\brelease management\b",
    r"\bplease let me know if you have any questions\b",
    r"\breport issues\b",
    r"\binformational\b",
]
SECRECY_PATTERNS = [
    r"\bconfidential\b",
    r"\bdo not discuss\b",
    r"\bkeep (?:this )?private\b",
    r"\bunreachable by phone\b",
]
DELIVERY_PATTERNS = [
    r"\bpackage\b",
    r"\bparcel\b",
    r"\bshipping\b",
    r"\bdelivery\b",
    r"\bdepot\b",
]
SENSITIVE_BRANDS = [
    r"\bmicrosoft 365\b",
    r"\bteams\b",
    r"\bvpn\b",
    r"\bpayroll\b",
    r"\bhr portal\b",
    r"\bdocu(?:sign)?\b",
]


def normalize_text(text: str) -> str:
    return WHITESPACE_RE.sub(" ", text.replace("\r", "\n")).strip()


def lower_text(text: str) -> str:
    return normalize_text(text).lower()


def count_pattern_hits(text: str, patterns: Iterable[str]) -> int:
    lowered = lower_text(text)
    return sum(len(re.findall(pattern, lowered, flags=re.IGNORECASE)) for pattern in patterns)


def extract_policy_signals(raw_email: str, intent_string: str) -> tuple[float, list[str]]:
    text = f"{raw_email}\n{intent_string}"
    lowered = lower_text(text)
    credential_hits = count_pattern_hits(lowered, CREDENTIAL_PATTERNS)
    payment_hits = count_pattern_hits(lowered, PAYMENT_PATTERNS)
    urgency_hits = count_pattern_hits(lowered, URGENCY_PATTERNS)
    authority_hits = count_pattern_hits(lowered, AUTHORITY_PATTERNS)
    document_hits = count_pattern_hits(lowered, DOCUMENT_PATTERNS)
    benign_hits = count_pattern_hits(lowered, BENIGN_PATTERNS)
    secrecy_hits = count_pattern_hits(lowered, SECRECY_PATTERNS)
    delivery_hits = count_pattern_hits(lowered, DELIVERY_PATTERNS)
    brand_hits = count_pattern_hits(lowered, SENSITIVE_BRANDS)

    credential_flow = (
        (
            credential_hits >= 2
            and (urgency_hits >= 1 or authority_hits >= 1 or brand_hits >= 1 or "locked out" in lowered)
        )
        or re.search(r"\b(?:sign in|log in|validate|verify).{0,40}\b(?:portal|vpn|account|mailbox|teams|microsoft 365)\b", lowered) is not None
    )
    payment_redirect = (
        payment_hits >= 2
        and re.search(r"\b(?:new|update|updated|change|changed|redirect)\b", lowered) is not None
        and re.search(r"\b(?:bank|beneficiary|wire|invoice|direct deposit|remittance)\b", lowered) is not None
    )
    sensitive_document = (
        document_hits >= 2
        and re.search(r"\b(?:compensation|legal|filing|contract|docusign|shared document)\b", lowered) is not None
        and (urgency_hits >= 1 or credential_hits >= 1)
    )
    delivery_release = (
        delivery_hits >= 1
        and re.search(r"\b(?:verify|confirm|payment|card|release)\b", lowered) is not None
    )
    routine_benign = (
        benign_hits >= 1
        and credential_hits == 0
        and secrecy_hits == 0
        and not credential_flow
        and not payment_redirect
        and not sensitive_document
        and not delivery_release
    )

    flags: list[str] = []
    probability = 0.0
    if credential_flow:
        flags.append("credential_flow")
        probability = max(probability, 0.92)
    if payment_redirect:
        flags.append("payment_redirect")
        probability = max(probability, 0.95)
    if sensitive_document:
        flags.append("sensitive_document")
        probability = max(probability, 0.88)
    if delivery_release:
        flags.append("delivery_release")
        probability = max(probability, 0.9)

    if probability == 0.0 and routine_benign:
        flags.append("routine_benign")
        probability = 0.08

    return probability, flags


def build_risk_feature_matrix(raw_emails: list[str], intent_strings: list[str]) -> np.ndarray:
    rows: list[list[float]] = []
    for raw_email, intent_string in zip(raw_emails, intent_strings, strict=False):
        normalized_email = normalize_text(raw_email)
        normalized_intent = normalize_text(intent_string)
        policy_probability, _ = extract_policy_signals(normalized_email, normalized_intent)
        rows.append(
            [
                float(len(URL_RE.findall(normalized_email))),
                float(len(EMAIL_RE.findall(normalized_email))),
                float(count_pattern_hits(normalized_email, CREDENTIAL_PATTERNS) + count_pattern_hits(normalized_intent, CREDENTIAL_PATTERNS)),
                float(count_pattern_hits(normalized_email, PAYMENT_PATTERNS) + count_pattern_hits(normalized_intent, PAYMENT_PATTERNS)),
                float(count_pattern_hits(normalized_email, URGENCY_PATTERNS) + count_pattern_hits(normalized_intent, URGENCY_PATTERNS)),
                float(count_pattern_hits(normalized_email, AUTHORITY_PATTERNS) + count_pattern_hits(normalized_intent, AUTHORITY_PATTERNS)),
                float(count_pattern_hits(normalized_email, DOCUMENT_PATTERNS) + count_pattern_hits(normalized_intent, DOCUMENT_PATTERNS)),
                float(count_pattern_hits(normalized_email, BENIGN_PATTERNS) + count_pattern_hits(normalized_intent, BENIGN_PATTERNS)),
                float(count_pattern_hits(normalized_email, SECRECY_PATTERNS) + count_pattern_hits(normalized_intent, SECRECY_PATTERNS)),
                float(count_pattern_hits(normalized_email, DELIVERY_PATTERNS) + count_pattern_hits(normalized_intent, DELIVERY_PATTERNS)),
                float(count_pattern_hits(normalized_email, SENSITIVE_BRANDS) + count_pattern_hits(normalized_intent, SENSITIVE_BRANDS)),
                float(len(normalized_email.split())),
                float(len(normalized_intent.split())),
                float(normalized_email.count("!")),
                float(normalized_email.count("?")),
                float(policy_probability),
            ]
        )
    return np.asarray(rows, dtype=np.float32)


def compose_lexical_text(raw_email: str, intent_string: str) -> str:
    normalized_email = normalize_text(raw_email)
    normalized_intent = normalize_text(intent_string)
    _, flags = extract_policy_signals(normalized_email, normalized_intent)
    flag_tokens = " ".join(f"policy_{flag}" for flag in flags)
    return f"{normalized_email}\nRID intent {normalized_intent}\n{flag_tokens}".strip()


@dataclass
class HybridTextVectorizer:
    word_max_features: int = 16000
    char_max_features: int = 12000

    def __post_init__(self) -> None:
        self.word_vectorizer = TfidfVectorizer(
            analyzer="word",
            ngram_range=(1, 2),
            lowercase=True,
            min_df=1,
            max_features=self.word_max_features,
            sublinear_tf=True,
            strip_accents="unicode",
        )
        self.char_vectorizer = TfidfVectorizer(
            analyzer="char_wb",
            ngram_range=(3, 5),
            lowercase=True,
            min_df=1,
            max_features=self.char_max_features,
            sublinear_tf=True,
            strip_accents="unicode",
        )

    def fit(self, raw_emails: list[str], intent_strings: list[str]) -> "HybridTextVectorizer":
        texts = [compose_lexical_text(raw_email, intent_string) for raw_email, intent_string in zip(raw_emails, intent_strings, strict=False)]
        self.word_vectorizer.fit(texts)
        self.char_vectorizer.fit(texts)
        return self

    def transform(self, raw_emails: list[str], intent_strings: list[str]) -> sparse.csr_matrix:
        texts = [compose_lexical_text(raw_email, intent_string) for raw_email, intent_string in zip(raw_emails, intent_strings, strict=False)]
        word_matrix = self.word_vectorizer.transform(texts)
        char_matrix = self.char_vectorizer.transform(texts)
        return sparse.hstack([word_matrix, char_matrix], format="csr")

    def fit_transform(self, raw_emails: list[str], intent_strings: list[str]) -> sparse.csr_matrix:
        self.fit(raw_emails, intent_strings)
        return self.transform(raw_emails, intent_strings)