Upload main.py

main.py

@@ -4,7 +4,8 @@ from typing import List, Dict, Tuple
 
 import gradio as gr
 
-# Optional imports for email classifier (loaded lazily)
+# Optional imports for email classifier (loaded lazily).
+# Space still runs if these aren't available (pure lexical fallback).
 try:
     import torch
     from transformers import AutoTokenizer, AutoModelForSequenceClassification
@@ -16,16 +17,16 @@ except Exception:
 # =========================
 # Config (env-overridable)
 # =========================
-EMAIL_CLASSIFIER_ID = os.getenv("EMAIL_CLASSIFIER_ID", "your-username/mini-phish")  # swap to your HF model repo later
-EMAIL_BACKBONE_ID   = os.getenv("EMAIL_BACKBONE_ID", "microsoft/MiniLM-L6-H384-uncased")
-THRESHOLD_TAU       = float(os.getenv("THRESHOLD_TAU", "0.40"))
-MAX_SEQ_LEN         = int(os.getenv("MAX_SEQ_LEN", "320"))
-SUBJECT_TOKEN_BUDGET= int(os.getenv("SUBJECT_TOKEN_BUDGET", "64"))
-FUSION_EMAIL_W      = float(os.getenv("FUSION_EMAIL_W", "0.6"))
-FUSION_URL_W        = float(os.getenv("FUSION_URL_W", "0.4"))
-URL_OVERRIDE_HIGH   = float(os.getenv("URL_OVERRIDE_HIGH", "0.85"))
-URL_OVERRIDE_KW     = float(os.getenv("URL_OVERRIDE_KW", "0.70"))
-ALLOWLIST_SAFE_CAP  = float(os.getenv("ALLOWLIST_SAFE_CAP", "0.15"))
+EMAIL_CLASSIFIER_ID   = os.getenv("EMAIL_CLASSIFIER_ID", "your-username/mini-phish")  # <- swap to your HF repo when ready
+EMAIL_BACKBONE_ID     = os.getenv("EMAIL_BACKBONE_ID", "microsoft/MiniLM-L6-H384-uncased")
+THRESHOLD_TAU         = float(os.getenv("THRESHOLD_TAU", "0.40"))
+MAX_SEQ_LEN           = int(os.getenv("MAX_SEQ_LEN", "320"))
+SUBJECT_TOKEN_BUDGET  = int(os.getenv("SUBJECT_TOKEN_BUDGET", "64"))
+FUSION_EMAIL_W        = float(os.getenv("FUSION_EMAIL_W", "0.6"))
+FUSION_URL_W          = float(os.getenv("FUSION_URL_W", "0.4"))
+URL_OVERRIDE_HIGH     = float(os.getenv("URL_OVERRIDE_HIGH", "0.85"))
+URL_OVERRIDE_KW       = float(os.getenv("URL_OVERRIDE_KW", "0.70"))
+ALLOWLIST_SAFE_CAP    = float(os.getenv("ALLOWLIST_SAFE_CAP", "0.15"))
 
 # =========================
 # Simple data classes
@@ -40,21 +41,22 @@ class UrlResult:
 class EmailResult:
     p_email: float
     kw_hits: List[str]
+    strong_hits: List[str]  # subset of kw_hits considered strong
 
 # =========================
-# URL extraction & heuristics (replace with your existing pipeline)
+# URL extraction & heuristics (swap with your real URL model when ready)
 # =========================
 URL_REGEX = r'(?i)\b((?:https?://|www\.)[^\s<>")]+)'
 
-SUSPICIOUS_TLDS = {".xyz", ".top", ".click", ".link", ".ru", ".cn", ".country", ".gq", ".ga", ".ml", ".tk"}
+SUSPICIOUS_TLDS = {
+    ".xyz", ".top", ".click", ".link", ".ru", ".cn", ".country", ".gq", ".ga", ".ml", ".tk"
+}
 SHORTENERS = {"bit.ly","t.co","tinyurl.com","goo.gl","ow.ly","is.gd","cutt.ly","tiny.one","lnkd.in"}
 
 def extract_urls(text: str) -> List[str]:
     if not text: return []
     urls = re.findall(URL_REGEX, text)
-    # normalize
-    uniq = []
-    seen = set()
+    uniq, seen = [], set()
     for u in urls:
         u = u.strip().strip(').,;\'"')
         if u and u not in seen:
@@ -95,29 +97,41 @@ def score_urls(urls: List[str]) -> List[UrlResult]:
 _tokenizer = None
 _model = None
 
-LEXICAL_CUES = [
-    "verify your account","update your password","immediately","within 24 hours",
-    "suspended","unusual activity","confirm","login","click","invoice","payment",
-    "otp","one-time password","unlock","reactivate","restricted","authenticate",
-    "security alert","urgent","limited time"
+# Strong vs normal cues (lowercase)
+STRONG_CUES = [
+    "otp", "one-time password", "one time password", "cvv", "pin", "pan",
+    "password", "bank details", "netbanking", "debit card", "credit card",
+    "lottery", "jackpot", "prize", "reward", "winner", "you have won",
+    "send otp", "share otp", "confirm otp", "verify otp",
+    "account restricted", "reactivate account", "unlock your account"
 ]
 
+NORMAL_CUES = [
+    "verify your account", "update your password", "immediately",
+    "within 24 hours", "suspended", "unusual activity", "confirm",
+    "login", "click", "invoice", "payment", "security alert",
+    "urgent", "limited time"
+]
+
+LEXICAL_CUES = sorted(set(STRONG_CUES + NORMAL_CUES))
+
 def load_email_model() -> Tuple[object, object]:
+    """Try to load EMAIL_CLASSIFIER_ID; on failure, fall back to backbone with small head.
+       Apply dynamic int8 quantization for CPU if available."""
     global _tokenizer, _model
     if _tokenizer is not None and _model is not None:
         return _tokenizer, _model
 
     if AutoTokenizer is None or AutoModelForSequenceClassification is None or torch is None:
-        # environment without torch/transformers (Space will still boot)
-        return None, None
+        return None, None  # environment without torch/transformers
 
-    # Try the preferred classifier first
+    # Preferred classifier
     model_id = EMAIL_CLASSIFIER_ID
     try:
         _tokenizer = AutoTokenizer.from_pretrained(model_id)
         _model = AutoModelForSequenceClassification.from_pretrained(model_id)
     except Exception:
-        # Fallback: load backbone and attach a tiny random head
+        # Fallback: backbone + fresh 2-class head
         try:
             _tokenizer = AutoTokenizer.from_pretrained(EMAIL_BACKBONE_ID)
             _model = AutoModelForSequenceClassification.from_pretrained(
@@ -127,7 +141,7 @@ def load_email_model() -> Tuple[object, object]:
             _tokenizer, _model = None, None
             return None, None
 
-    # Dynamic quantization for CPU
+    # Dynamic quantization (CPU)
     try:
         _model.eval()
         _model.to("cpu")
@@ -146,62 +160,78 @@ def _truncate_for_budget(tokens_subject: List[int], tokens_body: List[int], max_
     return subj + body
 
 def score_email(subject: str, body: str) -> EmailResult:
+    """Return EmailResult with probability + hit lists.
+       Strong cues push higher risk even without a model (email-only scams)."""
     text = (subject or "") + "\n" + (body or "")
-    # lightweight lexical cues for reasons + kw_flag
-    hits = [c for c in LEXICAL_CUES if c in text.lower()]
+    low = text.lower()
+
+    strong_hits = [c for c in STRONG_CUES if c in low]
+    normal_hits = [c for c in NORMAL_CUES if c in low]
+    all_hits = sorted(set(strong_hits + normal_hits))
 
     tok, mdl = load_email_model()
     if tok is None or mdl is None:
-        # fallback purely lexical probability
-        base = 0.15 + 0.1 * len(hits)
-        return EmailResult(p_email=float(min(base, 0.99)), kw_hits=hits)
+        # Pure lexical fallback (no model available):
+        base = 0.10
+        p_email = base + 0.18 * len(strong_hits) + 0.07 * len(normal_hits)
+        p_email = float(max(0.01, min(0.99, p_email)))
+        return EmailResult(p_email=p_email, kw_hits=all_hits, strong_hits=strong_hits)
 
-    # tokenize with budget
+    # Model path (MiniLM or your classifier)
     encoded_subj = tok.encode(subject or "", add_special_tokens=False)
     encoded_body = tok.encode(body or "", add_special_tokens=False)
-    input_ids = _truncate_for_budget(encoded_subj, encoded_body, MAX_SEQ_LEN-2, SUBJECT_TOKEN_BUDGET)
+    input_ids = _truncate_for_budget(encoded_subj, encoded_body, MAX_SEQ_LEN - 2, SUBJECT_TOKEN_BUDGET)
     input_ids = [tok.cls_token_id] + input_ids + [tok.sep_token_id]
-    attn_mask = [1]*len(input_ids)
+    attn_mask = [1] * len(input_ids)
 
-    import torch
     ids = torch.tensor([input_ids], dtype=torch.long)
     mask = torch.tensor([attn_mask], dtype=torch.long)
+
     with torch.no_grad():
         out = mdl(input_ids=ids, attention_mask=mask)
         if hasattr(out, "logits"):
-            logits = out.logits[0].detach().cpu().numpy().tolist()
-            # assume label 1 = phishing (prob via softmax)
             import math
+            logits = out.logits[0].detach().cpu().numpy().tolist()
             exps = [math.exp(x) for x in logits]
-            p1 = exps[1] / (exps[0] + exps[1])
+            p1 = exps[1] / (exps[0] + exps[1])  # assume label 1 = phishing
             p_email = float(p1)
         else:
             p_email = 0.5
 
-    # small calibration nudge from lexical cues (kept light)
-    p_email = float(min(0.99, max(0.01, p_email + 0.03*len(hits))))
-    return EmailResult(p_email=p_email, kw_hits=hits)
+    # Nudge with cues: stronger boost for strong hits
+    p_email += 0.10 * len(strong_hits) + 0.03 * len(normal_hits)
+    p_email = float(max(0.01, min(0.99, p_email)))
+
+    return EmailResult(p_email=p_email, kw_hits=all_hits, strong_hits=strong_hits)
 
 # =========================
 # Fusion
 # =========================
 def fuse(email_res: EmailResult, url_results: List[UrlResult], allowlist_domains: List[str]) -> Dict:
     r_url_max = max([u.risk for u in url_results], default=0.0)
-    kw_flag = 1 if email_res.kw_hits else 0
+    no_urls = (len(url_results) == 0)
 
-    # Allowlist check: if any URL host in allowlist
+    # Allowlist check: if any URL host in allowlist (only matters when URLs exist)
     allowlist_hit = False
     for u in url_results:
         h = url_host(u.url)
-        if any(h.endswith(d.lower()) for d in allowlist_domains):
+        if any(h.endswith(d.strip().lower()) for d in allowlist_domains if d.strip()):
             allowlist_hit = True
             break
 
+    # Base fusion
     r_total = FUSION_EMAIL_W * email_res.p_email + FUSION_URL_W * r_url_max
 
+    # URL-driven overrides
+    kw_flag = 1 if email_res.kw_hits else 0
     if (r_url_max >= URL_OVERRIDE_HIGH) or (kw_flag and r_url_max >= URL_OVERRIDE_KW):
         r_total = max(r_total, 0.90)
 
+    # Email-only strong-cue override
+    if no_urls and len(email_res.strong_hits) > 0:
+        r_total = max(r_total, 0.85)
+
+    # Allowlist cap
     if allowlist_hit:
         r_total = min(r_total, ALLOWLIST_SAFE_CAP)
 
@@ -211,6 +241,8 @@ def fuse(email_res: EmailResult, url_results: List[UrlResult], allowlist_domains
         "R_url_max": round(r_url_max, 3),
         "R_total": round(r_total, 3),
         "kw_hits": email_res.kw_hits,
+        "strong_hits": email_res.strong_hits,
+        "no_urls": no_urls,
         "allowlist_hit": allowlist_hit,
         "verdict": verdict
     }
@@ -219,17 +251,19 @@ def fuse(email_res: EmailResult, url_results: List[UrlResult], allowlist_domains
 # Gradio UI
 # =========================
 with gr.Blocks(title="PhishingMail-Lab") as demo:
-    gr.Markdown("# 🧪 PhishingMail‑Lab\nFree‑tier friendly POC with email+URL fusion")
+    gr.Markdown("# 🧪 PhishingMail‑Lab\n**POC** — Free‑tier friendly hybrid (email + URL) with explainable cues.")
 
     with gr.Row():
         subject = gr.Textbox(label="Subject", placeholder="Subject: Important account update")
-    body = gr.Textbox(label="Email Body (paste text or HTML)", lines=10, placeholder="Paste the email content here...")
+    body = gr.Textbox(label="Email Body (paste text or HTML)", lines=12, placeholder="Paste the email content here...")
     with gr.Row():
-        allowlist = gr.Textbox(label="Allowlist domains (comma-separated)", placeholder="microsoft.com, amazon.com")
+        allowlist = gr.Textbox(label="Allowlist domains (comma-separated)", placeholder="microsoft.com, amazon.in")
         tau = gr.Slider(0, 1, value=THRESHOLD_TAU, step=0.01, label="Decision Threshold τ")
     analyze_btn = gr.Button("Analyze")
 
     verdict = gr.Label(label="Verdict")
+    # NEW: context banner right under verdict
+    context_banner = gr.Markdown(visible=False)
     fusion_json = gr.JSON(label="Fusion & Flags")
     url_table = gr.Dataframe(headers=["URL","Risk","Reasons"], label="Per‑URL risk (heuristics demo)", interactive=False)
 
@@ -237,17 +271,27 @@ with gr.Blocks(title="PhishingMail-Lab") as demo:
         global THRESHOLD_TAU
         THRESHOLD_TAU = float(tau_val)
 
-        urls = list(dict.fromkeys(extract_urls((subject_text or "") + "\n" + (body_text or ""))))  # uniq while preserving order
+        # Extract URLs from both subject and body (keeps it simple)
+        urls = list(dict.fromkeys(extract_urls((subject_text or "") + "\n" + (body_text or ""))))  # uniq & ordered
         url_results = score_urls(urls)
         allow_domains = [d.strip().lower() for d in (allowlist_text or "").split(",") if d.strip()]
 
         email_res = score_email(subject_text or "", body_text or "")
         fused = fuse(email_res, url_results, allow_domains)
 
+        # Build banner text/visibility
+        banners = []
+        if fused.get("no_urls"):
+            banners.append("⚠️ **No URLs found** — decision based **only on email body**.")
+        if fused.get("allowlist_hit"):
+            banners.append("🛈 **Allowlist active** — risk **capped** for trusted domain.")
+        banner_text = "<br>".join(banners) if banners else ""
+        banner_visible = bool(banners)
+
         rows = [[u.url, round(u.risk,3), ", ".join(u.reasons)] for u in url_results]
-        return fused["verdict"], fused, rows
+        return fused["verdict"], gr.update(value=banner_text, visible=banner_visible), fused, rows
 
-    analyze_btn.click(run, [subject, body, allowlist, tau], [verdict, fusion_json, url_table])
+    analyze_btn.click(run, [subject, body, allowlist, tau], [verdict, context_banner, fusion_json, url_table])
 
 if __name__ == "__main__":
     demo.launch()