Update main.py

main.py

@@ -1,20 +1,32 @@
-import os, re, time, json
+import os, re, time, json, urllib.parse
 import gradio as gr
 import torch
 import torch.nn.functional as F
+import tldextract  # for robust registered-domain parsing
 
-# Be quiet + CPU friendly
+# Quiet + CPU friendly
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 
+# -------- Models / Regex --------
 URL_MODEL_ID = "CrabInHoney/urlbert-tiny-v4-malicious-url-classifier"
 LABEL_MAP = {0: "benign", 1: "defacement", 2: "malware", 3: "phishing"}
-#URL_RE = re.compile(r"""(?xi)\b(?:https?://|www\.)[a-z0-9\-._~%]+(?:/[^\s<>"']*)?""")
-URL_RE = re.compile(r"""(?xi)\bhttps?://[^\s<>"']+""")
+URL_RE = re.compile(r"""(?xi)\b(?:https?://|www\.)[a-z0-9\-._~%]+(?:/[^\s<>"']*)?""")
 
+# Heuristic config
+KEYWORDS = {
+    "login","verify","account","secure","update","bank","wallet",
+    "password","invoice","pay","reset","support","unlock","confirm"
+}
+SUSPICIOUS_TLDS = {
+    "zip","mov","lol","xyz","top","country","link","click","cam",
+    "help","gq","cf","tk","work","rest","monster","quest","live"
+}
 
+# Lazy globals for tokenizer & model
 _tok = None
 _mdl = None
 
+# -------- Utilities --------
 def _extract_urls(text: str):
     return sorted(set(m.group(0) for m in URL_RE.finditer(text or "")))
 
@@ -42,42 +54,125 @@ def _format_scores_md(scores_sorted):
         lines.append(f"| **{s['label']}** | {s['prob']*100:.2f} | {s['logit']:.3f} |")
     return "\n".join(lines)
 
-def _markdown_table(rows):
-    lines = ["| URL | Prediction | Confidence (%) |", "|---|---|---|"]
-    for u, lbl, conf in rows:
-        lines.append(f"| `{u}` | **{lbl}** | {conf:.2f} |")
+def _markdown_results_header(rows):
+    # rows: [ [url, model_label, model_pct, heur, fused, reason_txt], ... ]
+    lines = [
+        "| URL | Model | Model Prob (%) | Heuristic | Fused Risk | Reasons |",
+        "|---|---|---:|---:|---:|---|",
+    ]
+    for u, lbl, pct, h, fused, reasons in rows:
+        lines.append(
+            f"| `{u}` | **{lbl}** | {pct:.2f} | {h:.2f} | {fused:.2f} | {reasons} |"
+        )
     return "\n".join(lines)
 
 def _forensic_block(url, token_ids, tokens, scores_sorted, cls_vec, elapsed_s, truncated):
-    toks_preview = ", ".join(tokens[:64]) + (" …" if len(tokens) > 64 else "")
-    ids_preview  = ", ".join(map(str, token_ids[:64])) + (" …" if len(token_ids) > 64 else "")
+    toks_prev = ", ".join(tokens[:64]) + (" …" if len(tokens) > 64 else "")
+    ids_prev  = ", ".join(map(str, token_ids[:64])) + (" …" if len(token_ids) > 64 else "")
     cls_dim = len(cls_vec)
-    cls_preview = ", ".join(f"{v:.4f}" for v in cls_vec[:16]) + (" …" if cls_dim > 16 else "")
+    cls_prev = ", ".join(f"{v:.4f}" for v in cls_vec[:16]) + (" …" if cls_dim > 16 else "")
     l2 = (sum(v*v for v in cls_vec)) ** 0.5
-
     md = []
-    md.append(f"### 🔍 Forensics for `{url}`")
-    md.append("")
+    md.append(f"### 🔍 Forensics for `{url}`\n")
     md.append(f"- tokens: **{len(tokens)}** • truncated: **{'yes' if truncated else 'no'}**")
-    md.append(f"- inference time: **{elapsed_s:.2f}s**")
-    md.append("")
+    md.append(f"- inference time: **{elapsed_s:.2f}s**\n")
     md.append("**Top-k scores**")
     md.append(_format_scores_md(scores_sorted))
-    md.append("")
-    md.append("**Token IDs (preview)**")
-    md.append("```txt\n" + ids_preview + "\n```")
+    md.append("\n**Token IDs (preview)**")
+    md.append("```txt\n" + ids_prev + "\n```")
     md.append("**Tokens (preview)**")
-    md.append("```txt\n" + toks_preview + "\n```")
+    md.append("```txt\n" + toks_prev + "\n```")
     md.append("**[CLS] embedding (preview)**")
     md.append(f"`dim={cls_dim}`, `L2={l2:.4f}`")
-    md.append("```txt\n" + cls_preview + "\n```")
+    md.append("```txt\n" + cls_prev + "\n```")
     return "\n".join(md)
 
+# -------- Heuristics --------
+def _safe_parse(url: str):
+    # add scheme if missing so urlparse sees netloc
+    if not re.match(r"^https?://", url, re.I):
+        url = "http://" + url
+    return urllib.parse.urlparse(url)
+
+def heuristic_features(u: str):
+    feats = {}
+    try:
+        p = _safe_parse(u)
+        feats["scheme_https"] = 1 if p.scheme.lower() == "https" else 0
+        feats["host"] = p.hostname or ""
+        feats["path"] = p.path or "/"
+        feats["query"] = p.query or ""
+        ext = tldextract.extract(feats["host"])  # subdomain, domain, suffix
+        feats["registered_domain"] = f"{ext.domain}.{ext.suffix}" if ext.domain and ext.suffix else feats["host"]
+        feats["subdomain"] = ext.subdomain or ""
+        feats["tld"] = ext.suffix or ""
+        feats["labels"] = feats["host"].count(".") + (1 if feats["host"] else 0)
+        feats["has_at"] = "@" in u
+        feats["has_port"] = bool(p.netloc and ":" in p.netloc.split("@")[-1])
+        feats["has_punycode"] = "xn--" in feats["host"]
+        feats["len_url"] = len(u)
+        feats["hyphen_in_regdom"] = "-" in (ext.domain or "")
+        low_host = feats["host"].lower()
+        low_path = feats["path"].lower()
+        feats["kw_in_path"] = int(any(k in low_path for k in KEYWORDS))
+        feats["kw_in_host"] = int(any(k in low_host for k in KEYWORDS))
+        # keyword appears in subdomain but not in registered brand
+        feats["kw_in_subdomain_only"] = int(
+            feats["kw_in_host"] and (ext.domain and not any(k in ext.domain.lower() for k in KEYWORDS))
+        )
+        feats["suspicious_tld"] = int((feats["tld"].split(".")[-1] or "") in SUSPICIOUS_TLDS)
+        # crude “entropy-like” signal for long alnum query blobs
+        alnum = sum(c.isalnum() for c in feats["query"])
+        feats["query_ratio_alnum"] = (alnum / max(1, len(feats["query"]))) if feats["query"] else 0.0
+        feats["parse_error"] = False
+    except Exception:
+        feats = {"parse_error": True}
+    return feats
+
+def heuristic_score(feats: dict) -> float:
+    """0..1 suspicious score."""
+    if feats.get("parse_error"):
+        return 0.70  # unparsable => suspicious
+    score = 0.0
+    score += 0.25 * feats["kw_in_path"]
+    score += 0.20 * feats["kw_in_subdomain_only"]
+    score += 0.10 * feats["kw_in_host"]
+    score += 0.10 * feats["hyphen_in_regdom"]
+    score += 0.10 * (feats["labels"] >= 4)
+    score += 0.10 * feats["has_punycode"]
+    score += 0.10 * feats["suspicious_tld"]
+    score += 0.05 * feats["has_at"]
+    score += 0.05 * feats["has_port"]
+    score += 0.10 * (feats["len_url"] >= 100)
+    if feats["query"] and len(feats["query"]) >= 40 and feats["query_ratio_alnum"] > 0.9:
+        score += 0.10
+    return max(0.0, min(1.0, score))
+
+def heuristic_reasons(feats: dict) -> str:
+    if feats.get("parse_error"):
+        return "parse error"
+    rs = []
+    if feats.get("kw_in_path"): rs.append("keyword in path")
+    if feats.get("kw_in_subdomain_only"): rs.append("keyword in subdomain")
+    if feats.get("kw_in_host") and not feats.get("kw_in_subdomain_only"): rs.append("keyword in host")
+    if feats.get("hyphen_in_regdom"): rs.append("hyphen in registered domain")
+    if feats.get("labels", 0) >= 4: rs.append("deep subdomain nesting")
+    if feats.get("has_punycode"): rs.append("punycode host")
+    if feats.get("suspicious_tld"): rs.append(f"suspicious TLD: {feats.get('tld')}")
+    if feats.get("has_at"): rs.append("@ in URL")
+    if feats.get("has_port"): rs.append("explicit port")
+    if feats.get("len_url", 0) >= 100: rs.append("very long URL")
+    if feats.get("query") and len(feats.get("query", "")) >= 40 and feats.get("query_ratio_alnum", 0) > 0.9:
+        rs.append("long query blob")
+    return ", ".join(rs) if rs else "no heuristic triggers"
+
+# -------- Core --------
 def analyze(text: str, forensic: bool, show_json: bool):
     """
-    Returns a single Markdown block:
-      - verdict + compact table
-      - optional forensic blocks (tokens, logits, CLS)
+    One output: Markdown with
+      - verdict
+      - table (model, heuristic, fused + reasons)
+      - optional forensic blocks (tokens, logits, [CLS])
       - optional raw JSON (copy/paste)
     """
     text = (text or "").strip()
@@ -99,21 +194,15 @@ def analyze(text: str, forensic: bool, show_json: bool):
                 idx = int(k.split("_")[-1])
                 id2label[idx] = v
 
-    rows = []
-    unsafe = False
+    header_rows = []
     forensic_blocks = []
     export_data = {"model_id": URL_MODEL_ID, "items": []}
+    any_unsafe = False
 
     for u in urls:
+        # --- Encode & forward for logits / CLS ---
         max_len = min(512, getattr(mdl.config, "max_position_embeddings", 512) or 512)
-        enc = tok(
-            u,
-            truncation=True,
-            max_length=max_len,
-            padding=False,
-            return_tensors="pt",
-            return_attention_mask=True,
-        )
+        enc = tok(u, truncation=True, max_length=max_len, return_tensors="pt", return_attention_mask=True)
         token_ids = enc["input_ids"][0].tolist()
         tokens = tok.convert_ids_to_tokens(enc["input_ids"][0])
         truncated = enc["input_ids"].shape[1] >= max_len and len(tokens) >= max_len
@@ -123,9 +212,9 @@ def analyze(text: str, forensic: bool, show_json: bool):
             out = mdl(**enc, output_hidden_states=True)
         elapsed = time.time() - t0
 
-        logits = out.logits.squeeze(0)             # (num_labels,)
-        probs = _softmax(logits)                    # list[float]
-        hidden_states = out.hidden_states           # tuple of layers
+        logits = out.logits.squeeze(0)              # (num_labels,)
+        probs  = _softmax(logits)                   # list[float]
+        hidden_states = out.hidden_states
         cls_vec = hidden_states[-1][0, 0, :].cpu().tolist()
 
         per_class = [
@@ -134,10 +223,19 @@ def analyze(text: str, forensic: bool, show_json: bool):
         ]
         per_class_sorted = sorted(per_class, key=lambda x: x["prob"], reverse=True)
         top = per_class_sorted[0]
-        rows.append([u, top["label"], top["prob"] * 100.0])
-        if top["label"].lower() in {"phishing", "malware", "defacement"}:
-            unsafe = True
 
+        # --- Heuristics & fusion ---
+        feats = heuristic_features(u)
+        h_score = heuristic_score(feats)
+        mdl_phish_like = sum(s["prob"] for s in per_class_sorted if s["label"].lower() in {"phishing","malware","defacement"})
+        fused = 0.65 * mdl_phish_like + 0.35 * h_score
+        reasons = heuristic_reasons(feats)
+
+        header_rows.append([u, top["label"], top["prob"] * 100.0, h_score, fused, reasons])
+        if fused >= 0.50:
+            any_unsafe = True
+
+        # collect full details for optional JSON dump
         export_data["items"].append({
             "url": u,
             "token_ids": token_ids,
@@ -145,10 +243,13 @@ def analyze(text: str, forensic: bool, show_json: bool):
             "truncated": truncated,
             "logits": [float(x) for x in logits.cpu().tolist()],
             "probs":  [float(p) for p in probs],
-            "scores_sorted": per_class_sorted,  # label+prob+logit
+            "scores_sorted": per_class_sorted,
             "cls_vector": cls_vec,
             "cls_dim": len(cls_vec),
             "elapsed_sec": elapsed,
+            "heuristic": feats,
+            "heuristic_score": h_score,
+            "fused_risk": fused,
         })
 
         if forensic:
@@ -164,20 +265,20 @@ def analyze(text: str, forensic: bool, show_json: bool):
                 )
             )
 
-    verdict = "🔴 **UNSAFE (links flagged)**" if unsafe else "🟢 **SAFE (all links benign)**"
-    body = verdict + "\n\n" + _markdown_table(rows)
+    verdict = "🔴 **UNSAFE (links flagged)**" if any_unsafe else "🟢 **SAFE (no fused risk ≥ 0.50)**"
+    body = verdict + "\n\n" + _markdown_results_header(header_rows)
 
     if forensic and forensic_blocks:
         body += "\n\n---\n\n" + "\n\n---\n\n".join(forensic_blocks)
 
     if show_json:
-        # raw JSON for copy-paste (no File component needed)
         pretty = json.dumps(export_data, ensure_ascii=False, indent=2)
         body += "\n\n---\n\n**Raw forensics JSON (copy & save):**\n"
         body += "```json\n" + pretty + "\n```"
 
     return body
 
+# -------- UI --------
 demo = gr.Interface(
     fn=analyze,
     inputs=[
@@ -186,10 +287,13 @@ demo = gr.Interface(
         gr.Checkbox(label="Show raw JSON at the end (copy/paste)", value=False),
     ],
     outputs=gr.Markdown(label="Results"),
-    title="🛡️ PhishingMail — Forensics (HF Free CPU)",
-    description="Extract links, classify with a tiny URL model, and (optionally) view tokens, logits, and [CLS] embedding.",
+    title="🛡️ PhishingMail — Model + Heuristics (HF Free CPU)",
+    description=(
+        "We extract links and classify each with a compact malicious-URL model, then fuse with transparent heuristics. "
+        "Table shows Model Prob, Heuristic Score, and Fused Risk with reasons. "
+        "Toggle Forensic mode for tokens/logits/[CLS]."
+    ),
 )
 
 if __name__ == "__main__":
-    # Safe defaults for HF Spaces (no share=True needed)
     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)