Spaces:
Sleeping
Sleeping
| """ | |
| 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 | |