app.py

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)