import os import re import json from fastapi import FastAPI, HTTPException from pydantic import BaseModel from typing import Dict, Any, Optional, List import uvicorn from email import message_from_string from email.utils import parseaddr # --- ML Imports with Fallbacks --- try: import joblib import torch from transformers import AutoTokenizer, AutoModelForSequenceClassification from peft import PeftModel HAS_ML = True except ImportError as e: HAS_ML = False ML_ERROR = str(e) try: import pandas as pd from lime.lime_text import LimeTextExplainer import numpy as np HAS_LIME = True except ImportError: HAS_LIME = False # --- Ensemble Classifier Class --- class EnsembleClassifier: def __init__(self, models_dir=None): if models_dir is None: models_dir = os.path.join(os.path.dirname(__file__), "models") self.models_dir = models_dir self.header_rf = None self.attachment_rf = None # Text and URL transformers self.text_tokenizer = None self.text_model = None self.url_tokenizer = None self.url_model = None self.loaded_streams = { "text": False, "url": False, "header": False, "attachment": False } if not HAS_ML: raise ImportError(f"Missing machine learning dependencies: {ML_ERROR}") self.load_models() def load_models(self): # 1. Load Tabular Models (Header and Attachment Random Forests) header_path = os.path.join(self.models_dir, "header_rf.joblib") attachment_path = os.path.join(self.models_dir, "attachment_rf.joblib") if os.path.exists(header_path): self.header_rf = joblib.load(header_path) self.loaded_streams["header"] = True print("Loaded Header Random Forest model.") if os.path.exists(attachment_path): self.attachment_rf = joblib.load(attachment_path) self.loaded_streams["attachment"] = True print("Loaded Attachment Random Forest model.") # 2. Load URL model (downloads from Hugging Face Hub on first load) try: print("Loading pre-trained URL DistilBERT classifier (kmack/malicious-url-detection)...") self.url_tokenizer = AutoTokenizer.from_pretrained("kmack/malicious-url-detection") self.url_model = AutoModelForSequenceClassification.from_pretrained("kmack/malicious-url-detection") self.loaded_streams["url"] = True print("Loaded URL DistilBERT model.") except Exception as e: print(f"Warning: Could not load URL model: {e}") # 3. Load Text model (LoRA adapter + DistilBERT base) text_lora_dir = os.path.join(self.models_dir, "text_distilbert_lora") if os.path.exists(text_lora_dir): try: print("Loading fine-tuned Text DistilBERT model with LoRA...") self.text_tokenizer = AutoTokenizer.from_pretrained(text_lora_dir) base_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2) self.text_model = PeftModel.from_pretrained(base_model, text_lora_dir) self.loaded_streams["text"] = True print("Loaded Text LoRA model.") except Exception as e: print(f"Warning: Could not load Text model: {e}") else: print(f"Text model weights not found at {text_lora_dir}. Run train_text.py to generate them.") def predict_text_proba(self, text): if not self.loaded_streams["text"]: raise RuntimeError("Text classification model is not loaded.") inputs = self.text_tokenizer(text, truncation=True, padding=True, max_length=512, return_tensors="pt") with torch.no_grad(): outputs = self.text_model(**inputs) probs = torch.softmax(outputs.logits, dim=-1).numpy()[0] return probs.tolist() def predict_url_proba(self, url): if not self.loaded_streams["url"]: raise RuntimeError("URL classification model is not loaded.") inputs = self.url_tokenizer(url, truncation=True, padding=True, max_length=512, return_tensors="pt") with torch.no_grad(): outputs = self.url_model(**inputs) probs = torch.softmax(outputs.logits, dim=-1).numpy()[0] return probs.tolist() def predict(self, text: str, url: str, headers_features: dict, attachment_features: dict): is_text_active = bool(text and text.strip() and text != "Subject: \n\n" and self.loaded_streams["text"]) is_url_active = bool(url and url.strip() and self.loaded_streams["url"]) is_header_active = bool(headers_features and (headers_features.get("hdr_sender_name_len", 0) > 0 or headers_features.get("hdr_subject_len", 0) > 0) and self.loaded_streams["header"]) is_attachment_active = bool(attachment_features and attachment_features.get("att_has_attachment", 0) == 1 and self.loaded_streams["attachment"]) probabilities = {} verdicts = {} if is_text_active: probs_txt = self.predict_text_proba(text) probabilities["text"] = probs_txt verdicts["text"] = 1 if probs_txt[1] > 0.5 else 0 if is_url_active: probs_url = self.predict_url_proba(url) probabilities["url"] = probs_url verdicts["url"] = 1 if probs_url[1] > 0.5 else 0 if is_header_active: header_features_order = [ "hdr_sender_name_len", "hdr_sender_is_public", "hdr_is_spoofed", "hdr_subject_len", "hdr_subject_words", "hdr_has_urgency" ] features = pd.DataFrame([headers_features])[header_features_order] probs_hdr = self.header_rf.predict_proba(features)[0] probabilities["header"] = probs_hdr.tolist() verdicts["header"] = 1 if probs_hdr[1] > 0.5 else 0 if is_attachment_active: attachment_features_order = [ "att_has_attachment", "att_count", "att_avg_len", "att_is_dangerous", "att_has_double_ext", "att_is_sketchy_name" ] features = pd.DataFrame([attachment_features])[attachment_features_order] probs_att = self.attachment_rf.predict_proba(features)[0] probabilities["attachment"] = probs_att.tolist() verdicts["attachment"] = 1 if probs_att[1] > 0.5 else 0 active_verdicts = list(verdicts.values()) active_streams = list(verdicts.keys()) if not active_verdicts: return { "verdict": "Legitimate", "confidence": 1.0, "probabilities": {}, "active_streams": [] } if any(v == 1 for v in active_verdicts): final_verdict = "Phishing" phish_probs = [probabilities[stream][1] for stream in active_streams if verdicts[stream] == 1] confidence = max(phish_probs) if phish_probs else 0.5 else: final_verdict = "Legitimate" legit_probs = [probabilities[stream][0] for stream in active_streams] confidence = sum(legit_probs) / len(legit_probs) return { "verdict": final_verdict, "confidence": float(confidence), "probabilities": probabilities, "active_streams": active_streams } # --- Explainer Pipeline Class --- class ExplainerPipeline: def __init__(self, ensemble_clf): self.ensemble_clf = ensemble_clf self.lime_text_explainer = None if HAS_LIME: self.lime_text_explainer = LimeTextExplainer(class_names=["Legitimate", "Phishing"]) def explain_text_lime(self, text): if not HAS_LIME: raise ImportError("LIME library or PyTorch dependencies are not installed.") if not self.ensemble_clf.loaded_streams["text"]: raise RuntimeError("Text model is not loaded; cannot run LIME.") def classifier_predict_proba(texts): probs_list = [] for t in texts: probs = self.ensemble_clf.predict_text_proba(t) probs_list.append(probs) return np.array(probs_list) exp = self.lime_text_explainer.explain_instance( text, classifier_predict_proba, num_features=15, num_samples=100 ) return exp.as_list() def explain_tabular_perturbation(self, headers_features, attachment_features, result): explanations = { "headers": [], "attachments": [] } probs = result["probabilities"] # 1. Header Random Forest perturbation if "header" in probs and self.ensemble_clf.header_rf: base_phish_prob = probs["header"][1] benign_headers = { "hdr_sender_name_len": 5, "hdr_sender_is_public": 0, "hdr_is_spoofed": 0, "hdr_subject_len": 15, "hdr_subject_words": 3, "hdr_has_urgency": 0 } feature_descriptions = { "hdr_sender_name_len": "Sender display name length", "hdr_sender_is_public": "Public webmail domain sender", "hdr_is_spoofed": "Sender brand spoofing detected", "hdr_subject_len": "Subject line text length", "hdr_subject_words": "Subject word counts", "hdr_has_urgency": "Urgency indicator in subject" } for feature in headers_features.keys(): current_val = headers_features[feature] benign_val = benign_headers[feature] if current_val != benign_val: perturbed = headers_features.copy() perturbed[feature] = benign_val header_features_order = [ "hdr_sender_name_len", "hdr_sender_is_public", "hdr_is_spoofed", "hdr_subject_len", "hdr_subject_words", "hdr_has_urgency" ] df_perturbed = pd.DataFrame([perturbed])[header_features_order] perturbed_probs = self.ensemble_clf.header_rf.predict_proba(df_perturbed)[0] perturbed_phish_prob = perturbed_probs[1] weight = base_phish_prob - perturbed_phish_prob if abs(weight) > 0.01: explanations["headers"].append([ feature, float(weight), f"{feature_descriptions[feature]} (Value: {current_val} vs Benign: {benign_val})" ]) explanations["headers"].sort(key=lambda x: abs(x[1]), reverse=True) # 2. Attachment Random Forest perturbation if "attachment" in probs and self.ensemble_clf.attachment_rf: base_phish_prob = probs["attachment"][1] benign_attachments = { "att_has_attachment": 0, "att_count": 0, "att_avg_len": 0, "att_is_dangerous": 0, "att_has_double_ext": 0, "att_is_sketchy_name": 0 } feature_descriptions = { "att_has_attachment": "Email contains files", "att_count": "Attached file counts", "att_avg_len": "Filename character length", "att_is_dangerous": "Dangerous script/executable extension", "att_has_double_ext": "Double extension usage", "att_is_sketchy_name": "Sketchy name keywords" } for feature in attachment_features.keys(): current_val = attachment_features[feature] benign_val = benign_attachments[feature] if current_val != benign_val: perturbed = attachment_features.copy() perturbed[feature] = benign_val attachment_features_order = [ "att_has_attachment", "att_count", "att_avg_len", "att_is_dangerous", "att_has_double_ext", "att_is_sketchy_name" ] df_perturbed = pd.DataFrame([perturbed])[attachment_features_order] perturbed_probs = self.ensemble_clf.attachment_rf.predict_proba(df_perturbed)[0] perturbed_phish_prob = perturbed_probs[1] weight = base_phish_prob - perturbed_phish_prob if abs(weight) > 0.01: explanations["attachments"].append([ feature, float(weight), f"{feature_descriptions[feature]} (Value: {current_val} vs Benign: {benign_val})" ]) explanations["attachments"].sort(key=lambda x: abs(x[1]), reverse=True) return explanations def explain(self, text, url, headers_features, attachment_features, result): active_streams = result.get("active_streams", []) text_words = [] if "text" in active_streams and self.ensemble_clf.loaded_streams["text"]: text_words = self.explain_text_lime(text) tabular = self.explain_tabular_perturbation(headers_features, attachment_features, result) return { "verdict": result["verdict"], "confidence": result["confidence"], "text_words": text_words, "tabular": tabular, "is_lime_computed": HAS_LIME and self.ensemble_clf.loaded_streams["text"] } # --- FastAPI Initialization & Endpoints --- app = FastAPI( title="Multi-Modal Phishing & Social Engineering Detection API", description="REST API serving the real binary ensemble voting layer and LIME explanations.", version="1.2.0" ) # Relative OS-agnostic path resolution models_dir = os.path.join(os.path.dirname(__file__), "models") classifier = EnsembleClassifier(models_dir=models_dir) explainer = ExplainerPipeline(classifier) class ScanRequest(BaseModel): text: Optional[str] = None url: Optional[str] = None headers: Optional[Dict[str, Any]] = None attachments: Optional[Dict[str, Any]] = None raw_email: Optional[str] = None class ScanResponse(BaseModel): verdict: str confidence: float text_words: List[List[Any]] tabular: Dict[str, List[List[Any]]] is_lime_computed: bool probabilities: Dict[str, List[float]] active_streams: List[str] sender: Optional[str] = None subject: Optional[str] = None @app.get("/") def home(): return { "status": "online", "models_loaded": classifier.loaded_streams, "message": "Welcome to the real binary Multi-Modal Phishing API. Use /predict POST to scan or /metrics GET to fetch dataset performance metrics." } @app.get("/metrics") def get_metrics_endpoint(): metrics_path = os.path.join(models_dir, "metrics.json") if os.path.exists(metrics_path): try: with open(metrics_path, "r") as f: return json.load(f) except Exception: pass # If not cached, compute them test_path = os.path.join(os.path.dirname(__file__), "data", "processed", "test.csv") if not os.path.exists(test_path): raise HTTPException(status_code=404, detail=f"test.csv not found at {test_path}") try: from sklearn.metrics import precision_recall_fscore_support, confusion_matrix test_df = pd.read_csv(test_path).head(1000) def get_metrics_dict(y_true, y_pred): if not y_true or not y_pred: return { "precision": 0.0, "recall": 0.0, "f1": 0.0, "confusion_matrix": {"tn": 0, "fp": 0, "fn": 0, "tp": 0} } p, r, f, _ = precision_recall_fscore_support(y_true, y_pred, average="binary", zero_division=0) tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel() return { "precision": float(p), "recall": float(r), "f1": float(f), "confusion_matrix": { "tn": int(tn), "fp": int(fp), "fn": int(fn), "tp": int(tp) } } # 1. Header RF y_true_hdr = test_df["Label"].tolist() header_features = ["hdr_sender_name_len", "hdr_sender_is_public", "hdr_is_spoofed", "hdr_subject_len", "hdr_subject_words", "hdr_has_urgency"] df_hdr = test_df[header_features] y_pred_hdr = classifier.header_rf.predict(df_hdr).tolist() hdr_metrics = get_metrics_dict(y_true_hdr, y_pred_hdr) # 2. Attachment RF att_df = test_df[test_df["att_has_attachment"] == 1] if len(att_df) > 0: y_true_att = att_df["Label"].tolist() attachment_features = ["att_has_attachment", "att_count", "att_avg_len", "att_is_dangerous", "att_has_double_ext", "att_is_sketchy_name"] df_att = att_df[attachment_features] y_pred_att = classifier.attachment_rf.predict(df_att).tolist() att_metrics = get_metrics_dict(y_true_att, y_pred_att) else: att_metrics = get_metrics_dict([], []) # 3. URL Model url_df = test_df[test_df["URL"].notna() & (test_df["URL"].str.strip() != "")] y_pred_url_map = {} if classifier.loaded_streams["url"] and len(url_df) > 0: urls = url_df["URL"].tolist() y_true_url = url_df["Label"].tolist() device = "cuda" if torch.cuda.is_available() else "cpu" classifier.url_model.to(device) y_pred_url = [] batch_size = 64 for i in range(0, len(urls), batch_size): batch_urls = urls[i:i+batch_size] inputs = classifier.url_tokenizer(batch_urls, truncation=True, padding=True, max_length=128, return_tensors="pt").to(device) with torch.no_grad(): outputs = classifier.url_model(**inputs) probs = torch.softmax(outputs.logits, dim=-1).cpu().numpy() batch_preds = (probs[:, 1] > 0.5).astype(int).tolist() y_pred_url.extend(batch_preds) url_metrics = get_metrics_dict(y_true_url, y_pred_url) for idx, pred in zip(url_df.index, y_pred_url): y_pred_url_map[idx] = pred else: url_metrics = get_metrics_dict([], []) # 4. Ensemble y_true_ens = test_df["Label"].tolist() y_pred_ens = [] for idx, row in test_df.iterrows(): votes = [] votes.append(y_pred_hdr[idx]) if row["att_has_attachment"] == 1: feat_att = pd.DataFrame([row[attachment_features]]) att_vote = int(classifier.attachment_rf.predict(feat_att)[0]) votes.append(att_vote) if idx in y_pred_url_map: votes.append(y_pred_url_map[idx]) if any(v == 1 for v in votes): y_pred_ens.append(1) else: y_pred_ens.append(0) ens_metrics = get_metrics_dict(y_true_ens, y_pred_ens) res = { "header": hdr_metrics, "attachment": att_metrics, "url": url_metrics, "ensemble": ens_metrics, "dataset_info": { "total_records": len(test_df), "url_records": len(url_df), "attachment_records": len(att_df) } } try: with open(metrics_path, "w") as f: json.dump(res, f) except Exception: pass return res except Exception as e: raise HTTPException(status_code=500, detail=f"Metrics computation failed: {str(e)}") def parse_raw_email(raw_content: str) -> Dict[str, Any]: msg = message_from_string(raw_content) subject = msg.get("Subject", "") sender_raw = msg.get("From", "") display_name, email_address = parseaddr(sender_raw) name_len = len(display_name) domain = email_address.split('@')[1] if '@' in email_address else "" body = "" if msg.is_multipart(): for part in msg.walk(): content_type = part.get_content_type() content_disposition = str(part.get("Content-Disposition")) if content_type == "text/plain" and "attachment" not in content_disposition: try: body += part.get_payload(decode=True).decode(part.get_content_charset() or "utf-8", errors="ignore") except Exception: pass elif content_type == "text/html" and not body and "attachment" not in content_disposition: try: html_text = part.get_payload(decode=True).decode(part.get_content_charset() or "utf-8", errors="ignore") body += re.sub(r'<[^>]+>', ' ', html_text) except Exception: pass else: try: body = msg.get_payload(decode=True).decode(msg.get_content_charset() or "utf-8", errors="ignore") except Exception: pass urls = re.findall(r'https?://[^\s<>"]+|www\.[^\s<>"]+', body) url = urls[0] if urls else "" sender_is_public = 1 if domain.lower() in ["gmail.com", "yahoo.com", "hotmail.com", "outlook.com", "aol.com", "mail.com"] else 0 is_spoofed = 0 brands = ["dropbox", "pcloud", "microsoft", "google", "paypal", "dhl", "fedex", "netflix", "apple", "amazon", "citi", "bank"] dn_lower = display_name.lower() dom_lower = domain.lower() for brand in brands: if brand in dn_lower and brand not in dom_lower: is_spoofed = 1 break subject_cleaned = re.sub(r'\s+', ' ', subject).strip() subj_len = len(subject_cleaned) subj_words = len(subject_cleaned.split()) has_urgency = 0 urgency_words = ["urgent", "action required", "verify", "suspended", "security", "alert", "notice", "fail", "login", "password"] subject_lower = subject_cleaned.lower() for w in urgency_words: if w in subject_lower: has_urgency = 1 break headers_dict = { "hdr_sender_name_len": name_len, "hdr_sender_is_public": sender_is_public, "hdr_is_spoofed": is_spoofed, "hdr_subject_len": subj_len, "hdr_subject_words": subj_words, "hdr_has_urgency": has_urgency } has_attachment = 0 att_count = 0 att_avg_len = 0 att_is_dangerous = 0 att_has_double_ext = 0 att_is_sketchy_name = 0 filenames = [] dangerous_exts = {".exe", ".bat", ".scr", ".zip", ".rar", ".html", ".js", ".vbs", ".mobileconfig", ".jar", ".lnk", ".wsf"} sketchy_words = ["payment", "scan", "invoice", "approval", "bonuses", "calendar", "agreement", "contract", "card", "statement", "verify", "update", "doc", "pdf"] for part in msg.walk(): filename = part.get_filename() if filename: filenames.append(filename) att_count += 1 has_attachment = 1 f_lower = filename.lower() _, ext = os.path.splitext(f_lower) if ext in dangerous_exts: att_is_dangerous = 1 parts = f_lower.split('.') if len(parts) > 2: att_has_double_ext = 1 for w in sketchy_words: if w in f_lower: att_is_sketchy_name = 1 break if att_count > 0: att_avg_len = sum(len(f) for f in filenames) / att_count attachments_dict = { "att_has_attachment": has_attachment, "att_count": att_count, "att_avg_len": att_avg_len, "att_is_dangerous": att_is_dangerous, "att_has_double_ext": att_has_double_ext, "att_is_sketchy_name": att_is_sketchy_name } return { "text": f"Subject: {subject}\n\n{body}", "url": url, "headers": headers_dict, "attachments": attachments_dict, "sender": sender_raw, "subject": subject } @app.post("/predict", response_model=ScanResponse) def predict(request: ScanRequest): sender_val = None subject_val = None if request.raw_email: try: parsed = parse_raw_email(request.raw_email) text = parsed["text"] url = parsed["url"] headers = parsed["headers"] attachments = parsed["attachments"] sender_val = parsed["sender"] subject_val = parsed["subject"] except Exception as e: raise HTTPException(status_code=400, detail=f"Failed to parse raw email: {str(e)}") else: text = request.text url = request.url headers = request.headers attachments = request.attachments try: res = classifier.predict(text, url, headers, attachments) exp = explainer.explain(text, url, headers, attachments, res) except RuntimeError as e: raise HTTPException(status_code=500, detail=str(e)) return ScanResponse( verdict=exp["verdict"], confidence=float(exp["confidence"]), text_words=exp["text_words"], tabular=exp["tabular"], is_lime_computed=exp["is_lime_computed"], probabilities=res["probabilities"], active_streams=res["active_streams"], sender=sender_val, subject=subject_val ) if __name__ == "__main__": uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)