src/feature_engineering.py

"""
Feature engineering and NLP preprocessing for email threat classification.
Provides transformers for lexical, behavioral, and structural threat indicators.
"""

import re
import html
import numpy as np
import pandas as pd
from typing import List, Dict, Any, Union
from sklearn.base import BaseEstimator, TransformerMixin


# Regular expression patterns for threat detection
HTML_LINK_REGEX = re.compile(r'<a\s+[^>]*href=[\'"]([^\'"]+)[\'"][^>]*>(.*?)</a>', re.IGNORECASE | re.DOTALL)
URL_REGEX = re.compile(r'https?://[^\s<>"]+|www\.[^\s<>"]+', re.IGNORECASE)
IP_URL_REGEX = re.compile(r'https?://(?:\d{1,3}\.){3}\d{1,3}(?::\d+)?(?:/[^\s<>"]*)?', re.IGNORECASE)
SHORTENER_REGEX = re.compile(r'https?://(?:bit\.ly|tinyurl\.com|t\.co|goo\.gl|ow\.ly|is\.gd|buff\.ly|rebrand\.ly)/[^\s<>"]*', re.IGNORECASE)
EMAIL_REGEX = re.compile(r'[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+')
HTML_TAG_REGEX = re.compile(r'<[^>]+>')
PHONE_REGEX = re.compile(r'\b(?:\+?\d{1,3}[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b')

# Lexicon dictionaries for domain-specific indicators
URGENCY_KEYWORDS = [
    "urgent", "immediately", "immediate", "within 24 hours", "within 12 hours",
    "within 2 hours", "within 30 minutes", "within 48 hours",
    "act now", "suspended", "suspension", "terminated", "deadline", "expires today",
    "critical", "action required", "respond promptly", "final notice", "attention required",
    "time sensitive", "limited time", "don't delay", "freeze", "account locked"
]

CREDENTIAL_KEYWORDS = [
    "verify your account", "confirm password", "enter credentials", "login to continue",
    "password reset", "security update", "authenticate", "identity verification",
    "click here to login", "sign in below", "update credentials", "otp", "one-time password",
    "unlock account", "validate email", "re-activate", "access restricted", "verify your employee identity",
    "unauthorized login", "restore access"
]

FINANCIAL_KEYWORDS = [
    "wire transfer", "bank account", "routing number", "invoice attached",
    "overdue payment", "bitcoin", "crypto", "ethereum", "wallet address",
    "gift card", "remittance", "swift code", "direct deposit", "payroll update",
    "billing statement", "unauthorized transaction", "charge refund", "purchase order",
    "escrow account"
]

IMPERSONATION_KEYWORDS = [
    "ceo", "chief executive", "cfo", "director", "executive office", "board member",
    "it support", "it desk", "helpdesk", "system administrator", "admin team",
    "security operations", "compliance department", "human resources", "microsoft support",
    "office 365", "microsoft 365", "google workspace", "paypal", "docusign", "irs"
]

ATTACHMENT_KEYWORDS = [
    ".pdf", ".zip", ".exe", ".iso", ".scr", ".rar", ".docm", ".xlsm",
    "invoice.pdf", "receipt.zip", "statement.pdf", "attachment", "enclosed document"
]


def clean_text(text: str) -> str:
    """
    Cleans raw text while preserving linguistic threat cues.
    Decodes HTML entities, preserves URLs from link tags, strips HTML tags, and normalizes whitespace.
    """
    if not isinstance(text, str):
        return ""
    # Unescape HTML entities
    cleaned = html.unescape(text)
    # Convert <a href='...'>text</a> into text [URL] so URLs in links are not lost
    cleaned = HTML_LINK_REGEX.sub(r' \2 [URL] ', cleaned)
    # Replace remaining standalone URLs with normalized token
    cleaned = URL_REGEX.sub(" [URL] ", cleaned)
    # Remove remaining HTML tags
    cleaned = HTML_TAG_REGEX.sub(" ", cleaned)
    # Normalize whitespaces
    cleaned = re.sub(r'\s+', ' ', cleaned).strip()
    return cleaned


def extract_threat_indicators(subject: str, body: str) -> Dict[str, Any]:
    """
    Extracts structured cyber threat indicators and behavioral cues from an email.
    Used for both feature extraction and real-time UI explainability.
    """
    subject = str(subject or "")
    body = str(body or "")
    combined = f"{subject}\n{body}"
    combined_lower = combined.lower()

    # URL features
    all_urls = URL_REGEX.findall(combined)
    ip_urls = IP_URL_REGEX.findall(combined)
    shortened_urls = SHORTENER_REGEX.findall(combined)

    # Lexical keyword matching
    urgency_hits = [kw for kw in URGENCY_KEYWORDS if kw in combined_lower]
    credential_hits = [kw for kw in CREDENTIAL_KEYWORDS if kw in combined_lower]
    financial_hits = [kw for kw in FINANCIAL_KEYWORDS if kw in combined_lower]
    impersonation_hits = [kw for kw in IMPERSONATION_KEYWORDS if kw in combined_lower]
    attachment_hits = [kw for kw in ATTACHMENT_KEYWORDS if kw in combined_lower]

    # Structural metrics
    total_chars = len(combined)
    caps_count = sum(1 for c in combined if c.isupper())
    caps_ratio = (caps_count / total_chars) if total_chars > 0 else 0.0
    exclamation_count = combined.count("!")
    question_count = combined.count("?")
    dollar_count = combined.count("$") + combined.count("€") + combined.count("£")

    # Risk score heuristics (0.0 to 10.0 scale)
    heuristic_score = (
        len(urgency_hits) * 1.5 +
        len(credential_hits) * 2.0 +
        len(financial_hits) * 1.5 +
        len(impersonation_hits) * 1.2 +
        len(all_urls) * 0.8 +
        len(ip_urls) * 3.0 +
        len(shortened_urls) * 2.0 +
        (1.5 if caps_ratio > 0.25 else 0.0)
    )

    return {
        "url_count": len(all_urls),
        "ip_url_count": len(ip_urls),
        "shortened_url_count": len(shortened_urls),
        "urgency_hits": urgency_hits,
        "urgency_count": len(urgency_hits),
        "credential_hits": credential_hits,
        "credential_count": len(credential_hits),
        "financial_hits": financial_hits,
        "financial_count": len(financial_hits),
        "impersonation_hits": impersonation_hits,
        "impersonation_count": len(impersonation_hits),
        "attachment_hits": attachment_hits,
        "attachment_count": len(attachment_hits),
        "caps_ratio": round(caps_ratio, 4),
        "exclamation_count": exclamation_count,
        "question_count": question_count,
        "dollar_count": dollar_count,
        "body_length": len(body),
        "subject_length": len(subject),
        "heuristic_threat_score": round(min(heuristic_score, 10.0), 2)
    }


class ThreatFeatureExtractor(BaseEstimator, TransformerMixin):
    """
    Scikit-learn compatible transformer extracting numeric cybersecurity
    and behavioral features from email (subject, body) pairs.
    """

    FEATURE_NAMES = [
        "url_count",
        "ip_url_count",
        "shortened_url_count",
        "urgency_count",
        "credential_count",
        "financial_count",
        "impersonation_count",
        "attachment_count",
        "caps_ratio",
        "exclamation_count",
        "question_count",
        "dollar_count",
        "heuristic_score"
    ]

    def __init__(self):
        pass

    def fit(self, X, y=None):
        return self

    def transform(self, X) -> np.ndarray:
        """
        Transforms input DataFrame or records into feature array.
        Expects X to have 'subject' and 'body' columns/keys.
        """
        if isinstance(X, pd.DataFrame):
            records = X.to_dict(orient="records")
        elif isinstance(X, list):
            records = X
        else:
            raise ValueError("Input X must be a pandas DataFrame or list of dicts.")

        features = []
        for row in records:
            subj = str(row.get("subject", "") or "")
            body = str(row.get("body", "") or "")
            indicators = extract_threat_indicators(subj, body)
            
            row_features = [
                indicators["url_count"],
                indicators["ip_url_count"],
                indicators["shortened_url_count"],
                indicators["urgency_count"],
                indicators["credential_count"],
                indicators["financial_count"],
                indicators["impersonation_count"],
                indicators["attachment_count"],
                indicators["caps_ratio"],
                indicators["exclamation_count"],
                indicators["question_count"],
                indicators["dollar_count"],
                indicators["heuristic_threat_score"]
            ]
            features.append(row_features)

        return np.array(features, dtype=np.float32)

    def get_feature_names_out(self, input_features=None):
        return np.array(self.FEATURE_NAMES)


class EmailTextCombiner(BaseEstimator, TransformerMixin):
    """
    Combines 'subject' and 'body' into a single cleaned text string for TF-IDF vectorization.
    """

    def __init__(self, subject_weight: int = 2):
        self.subject_weight = subject_weight

    def fit(self, X, y=None):
        return self

    def transform(self, X) -> List[str]:
        if isinstance(X, pd.DataFrame):
            records = X.to_dict(orient="records")
        elif isinstance(X, list):
            records = X
        else:
            raise ValueError("Input X must be a pandas DataFrame or list of dicts.")

        combined_texts = []
        for row in records:
            subj = clean_text(str(row.get("subject", "") or ""))
            body = clean_text(str(row.get("body", "") or ""))
            
            # Repeat subject to amplify subject line keywords
            weighted_subj = " ".join([subj] * self.subject_weight) if subj else ""
            combined = f"{weighted_subj} {body}".strip()
            combined_texts.append(combined if combined else "empty_message")
            
        return combined_texts