Update the CyberForge ML service

trainer.py

@@ -1,6 +1,13 @@
 """
 Trains the CyberForge detection models from public, labelled data.
 
+(Second version. The first took its benign and phishing URLs from one older
+collection whose benign URLs are deep links from a crawl and whose phishing
+URLs are mostly short. A dry run showed the phishing model had learned exactly
+that, path length, and it scored bbc.co.uk/news at 0.96; the acceptance gate
+rejected it. The benign and phishing examples now come from several sources
+that overlap in shape, including the landing pages of popular sites.)
+
 This runs inside the Hugging Face Space (app.py starts it from a protected
 endpoint). It is never run on a developer machine.
 
@@ -47,19 +54,25 @@ from url_features import (ALL_FEATURES, ATTACK_FEATURES, LEXICAL_FEATURES, bare_
 
 logger = logging.getLogger("cyberforge.trainer")
 
-URL_COLLECTION = "bgspaditya/byt-malicious-url-treatment"   # ~640k URLs: benign / phishing / malware / defacement
+PHISHING_COLLECTION = "Mitake/PhishingURLsANDBenignURLs"     # ~886k URLs: benign (0) / phishing (1), CC-BY-SA-4.0
+URL_COLLECTION = "bgspaditya/byt-malicious-url-treatment"   # ~640k URLs; used for its malware URLs and some benign ones
 PHISHING_REFERENCE = "pirocheto/phishing-url"                 # 11k balanced full URLs, CC-BY-4.0, with its own test split
+POPULAR_SITES = "phl-ldm/tranco_top_million"                  # landing pages of top-million sites, CC-BY-4.0 (one small file of it)
 ATTACK_PAYLOADS = "shengqin/web-attacks"                      # SQL-injection and XSS payloads, and normal strings
 
 DEFAULTS = {
-    "benign_sample": 120000,
-    "phishing_sample": 80000,
+    "benign_sample": 120000,            # benign URLs from the phishing collection
+    "collection_benign_sample": 40000,  # benign URLs from the older collection (deep links)
+    "popular_sample": 60000,            # landing pages of popular sites
+    "phishing_sample": 120000,
     "malware_sample": 30000,
     "attack_benign_sample": 40000,
     "test_fraction": 0.2,
     "max_iter": 300,
     "learning_rate": 0.1,
     "max_leaf_nodes": 31,
+    "min_samples_leaf": 20,
+    "l2_regularization": 0.0,
     "anomaly_contamination": 0.02,
     "seed": 42,
 }
@@ -69,8 +82,10 @@ DEFAULTS = {
 # none of the well-known legitimate URLs below.
 GATE = {"min_auc": 0.90, "max_false_positive_rate": 0.05, "max_known_legitimate_flagged": 0}
 # The anomaly model has no labels to be scored against, so its gate is that it
-# flags few held-out benign URLs and more malicious ones than benign ones.
-ANOMALY_GATE = {"max_benign_flagged": 0.05}
+# flags few held-out benign URLs and clearly more malicious ones. (The first
+# version only asked for "more"; a model that flagged 1.8% of malicious URLs
+# against 1.6% of benign ones passed, which is no separation at all.)
+ANOMALY_GATE = {"max_benign_flagged": 0.05, "min_malicious_flagged": 0.10, "min_ratio": 3.0}
 
 # Well-known legitimate URLs, several of them long or full of parameters: the
 # kind the first models called critical. They are never used for training.
@@ -163,10 +178,29 @@ def clean_urls(values) -> List[str]:
     return out
 
 
+def load_phishing_collection(token: Optional[str]) -> Dict[str, List[str]]:
+    path = _download(PHISHING_COLLECTION, lambda name: name.endswith(".csv"), token)[0]
+    frame = pd.read_csv(path, usecols=["url", "label"])
+    return {"benign": clean_urls(frame.loc[frame["label"] == 0, "url"].tolist()),
+            "phishing": clean_urls(frame.loc[frame["label"] == 1, "url"].tolist())}
+
+
 def load_url_collection(token: Optional[str]) -> Dict[str, List[str]]:
     paths = _download(URL_COLLECTION, lambda name: name.startswith("data/train") and name.endswith(".parquet"), token)
     frame = pd.concat([pd.read_parquet(path, columns=["url", "type"]) for path in paths], ignore_index=True)
-    return {kind: clean_urls(frame.loc[frame["type"] == kind, "url"].tolist()) for kind in ("benign", "phishing", "malware")}
+    return {kind: clean_urls(frame.loc[frame["type"] == kind, "url"].tolist()) for kind in ("benign", "malware")}
+
+
+def load_popular_sites(token: Optional[str]) -> List[str]:
+    """Landing pages of popular sites: the page a visitor actually ends up on,
+    and the site's bare address. Only the dataset's small consent file is read."""
+    path = _download(POPULAR_SITES, lambda name: name.endswith("cookie_warning_consents.parquet"), token)[0]
+    frame = pd.read_parquet(path)
+    urls: List[str] = []
+    for column in ("url", "crawled_url"):
+        if column in frame.columns:
+            urls += frame[column].tolist()
+    return clean_urls(urls)
 
 
 def load_phishing_reference(token: Optional[str]) -> Dict[str, Dict[str, List[str]]]:
@@ -199,6 +233,18 @@ def load_payloads(token: Optional[str]) -> Dict[str, List[str]]:
     return out
 
 
+def labelled(parts: Dict[str, List[str]]) -> List[tuple]:
+    """[(url, source), ...] from {source: [urls]}, each URL once."""
+    seen, rows = set(), []
+    for source, urls in parts.items():
+        for url in urls:
+            key = bare_url(url).lower()
+            if key and key not in seen:
+                seen.add(key)
+                rows.append((url, source))
+    return rows
+
+
 def compose_url(base: str, payload: str, rng: random.Random) -> str:
     """Put a payload where one arrives in practice: as a query value on an
     ordinary site, encoded or not, or (less often) in the path."""
@@ -258,11 +304,20 @@ def score_urls(model, names: List[str], urls: List[str]) -> List[Dict]:
     return [{"url": url, "score": round(float(score), 4)} for url, score in zip(urls, scores)]
 
 
-def train_classifier(name: str, positives: List[str], negatives: List[str], names: List[str], settings: Dict,
-                     extra_tests: Optional[Dict[str, Dict[str, List[str]]]] = None) -> Dict:
-    """Fit one classifier and measure it on sites it did not train on."""
-    urls = negatives + positives
-    labels = np.array([0] * len(negatives) + [1] * len(positives))
+def train_classifier(name: str, positives: Dict[str, List[str]], negatives: Dict[str, List[str]], names: List[str],
+                     settings: Dict, extra_tests: Optional[Dict[str, Dict[str, List[str]]]] = None) -> Dict:
+    """Fit one classifier and measure it on sites it did not train on.
+    `positives` and `negatives` are {source: [urls]}; results are also broken
+    down by source, so one source cannot hide behind another."""
+    negative_rows, positive_rows = labelled(negatives), labelled(positives)
+    # a URL listed on both sides teaches nothing; drop it from both
+    clash = {bare_url(url).lower() for url, _ in negative_rows} & {bare_url(url).lower() for url, _ in positive_rows}
+    negative_rows = [row for row in negative_rows if bare_url(row[0]).lower() not in clash]
+    positive_rows = [row for row in positive_rows if bare_url(row[0]).lower() not in clash]
+
+    urls = [url for url, _ in negative_rows + positive_rows]
+    sources = np.array([source for _, source in negative_rows + positive_rows])
+    labels = np.array([0] * len(negative_rows) + [1] * len(positive_rows))
     groups = np.array([registered_domain(url) for url in urls])
     splitter = GroupShuffleSplit(n_splits=1, test_size=settings["test_fraction"], random_state=settings["seed"])
     train_index, test_index = next(splitter.split(urls, labels, groups))
@@ -270,14 +325,23 @@ def train_classifier(name: str, positives: List[str], negatives: List[str], name
     table = feature_table(urls)[:, columns(names)]
     model = HistGradientBoostingClassifier(
         max_iter=int(settings["max_iter"]), learning_rate=float(settings["learning_rate"]),
-        max_leaf_nodes=int(settings["max_leaf_nodes"]), class_weight="balanced",
+        max_leaf_nodes=int(settings["max_leaf_nodes"]), min_samples_leaf=int(settings["min_samples_leaf"]),
+        l2_regularization=float(settings["l2_regularization"]), class_weight="balanced",
         early_stopping=True, validation_fraction=0.1, n_iter_no_change=15, random_state=settings["seed"],
     )
     started = time.time()
     model.fit(table[train_index], labels[train_index])
     seconds = round(time.time() - started, 1)
 
-    holdout = measure(labels[test_index], model.predict_proba(table[test_index])[:, 1])
+    test_scores = model.predict_proba(table[test_index])[:, 1]
+    holdout = measure(labels[test_index], test_scores)
+    by_source = {}
+    for source in sorted(set(sources[test_index])):
+        mask = sources[test_index] == source
+        flagged = float(np.mean(test_scores[mask] >= 0.5))
+        positive = bool(labels[test_index][mask][0])
+        by_source[source] = {"urls": int(mask.sum()), "kind": "positive" if positive else "negative",
+                             ("recall" if positive else "false_positive_rate"): round(flagged, 4)}
     report = {
         "model_type": "HistGradientBoostingClassifier",
         "feature_names": names,
@@ -285,8 +349,10 @@ def train_classifier(name: str, positives: List[str], negatives: List[str], name
         "training": {
             "rows": int(len(train_index)), "positives": int(labels[train_index].sum()),
             "sites": int(len(set(groups[train_index]))), "seconds": seconds, "iterations": int(model.n_iter_),
+            "rows_by_source": {source: int(np.sum(sources[train_index] == source)) for source in sorted(set(sources))},
+            "dropped_as_listed_on_both_sides": len(clash),
         },
-        "holdout": {**holdout, "sites": int(len(set(groups[test_index]))),
+        "holdout": {**holdout, "sites": int(len(set(groups[test_index]))), "by_source": by_source,
                     "note": "URLs from sites that were not in the training rows"},
         "other_datasets": {},
     }
@@ -369,8 +435,8 @@ def train_anomaly(benign: List[str], malicious: List[str], names: List[str], set
     reasons = []
     if benign_flagged > ANOMALY_GATE["max_benign_flagged"]:
         reasons.append(f"it flags {benign_flagged} of held-out benign URLs, above {ANOMALY_GATE['max_benign_flagged']}")
-    if malicious_flagged <= benign_flagged:
-        reasons.append("it does not flag malicious URLs more often than benign ones")
+    if malicious_flagged < ANOMALY_GATE["min_malicious_flagged"] or malicious_flagged < ANOMALY_GATE["min_ratio"] * benign_flagged:
+        reasons.append(f"it flags {malicious_flagged} of malicious URLs against {benign_flagged} of benign ones, which is not a useful separation")
     report["gate"] = {"rules": ANOMALY_GATE, "passed": not reasons, "reasons": reasons}
     return {"model": model, "report": report}
 
@@ -389,15 +455,25 @@ def run_training(options: Dict, token: Optional[str], model_repo: str, progress:
     wanted = options.get("models") or ["phishing_detection", "malware_detection", "web_attack_detection", "anomaly_detection"]
     dry_run = bool(options.get("dry_run", False))
     rng = random.Random(settings["seed"])
+    # Features a run is asked to leave out, to measure whether a model depends
+    # on something it should not. The list is recorded with the model.
+    dropped = [name for name in (options.get("drop_features") or []) if isinstance(name, str)]
+    lexical = [name for name in LEXICAL_FEATURES if name not in dropped]
+    attack = [name for name in ATTACK_FEATURES if name not in dropped]
+    settings["drop_features"] = dropped
 
     progress("downloading data")
+    phishing = load_phishing_collection(token)
     collection = load_url_collection(token)
     reference = load_phishing_reference(token)
+    popular = load_popular_sites(token)
     payloads = load_payloads(token) if "web_attack_detection" in wanted else {"attack": [], "normal": []}
     sources = {
+        PHISHING_COLLECTION: {"revision": _revision(PHISHING_COLLECTION, token), "rows": {kind: len(urls) for kind, urls in phishing.items()}},
         URL_COLLECTION: {"revision": _revision(URL_COLLECTION, token), "rows": {kind: len(urls) for kind, urls in collection.items()}},
         PHISHING_REFERENCE: {"revision": _revision(PHISHING_REFERENCE, token),
                              "rows": {split: {status: len(urls) for status, urls in parts.items()} for split, parts in reference.items()}},
+        POPULAR_SITES: {"revision": _revision(POPULAR_SITES, token), "rows": {"landing pages": len(popular)}},
         ATTACK_PAYLOADS: {"revision": _revision(ATTACK_PAYLOADS, token), "rows": {kind: len(texts) for kind, texts in payloads.items()}},
     }
 
@@ -406,54 +482,66 @@ def run_training(options: Dict, token: Optional[str], model_repo: str, progress:
         rng.shuffle(items)
         return items[: int(size)]
 
-    benign = sample(collection["benign"], settings["benign_sample"])
     # URLs of the reference test split must not be trained on.
     reference_test = {bare_url(url).lower() for part in reference["test"].values() for url in part}
     keep = lambda urls: [url for url in urls if bare_url(url).lower() not in reference_test]  # noqa: E731
 
+    # Benign URLs of several shapes, so that no model can tell benign from
+    # hostile by how one list happens to be written: deep links, short
+    # addresses and bare domains, landing pages of popular sites.
+    benign = {
+        "phishing collection, benign": keep(sample(phishing["benign"], settings["benign_sample"])),
+        "older collection, benign": keep(sample(collection["benign"], settings["collection_benign_sample"])),
+        "popular sites": keep(sample(popular, settings["popular_sample"])),
+        "reference, legitimate": keep(reference["train"]["legitimate"]),
+    }
+    all_benign = [url for urls in benign.values() for url in urls]
+
     results: Dict[str, Dict] = {}
-    try_models = [name for name in wanted]
 
-    if "phishing_detection" in try_models:
+    if "phishing_detection" in wanted:
         progress("training phishing_detection")
-        positives = keep(sample(collection["phishing"], settings["phishing_sample"]) + reference["train"]["phishing"])
-        negatives = keep(benign + reference["train"]["legitimate"])
+        positives = {
+            "phishing collection, phishing": keep(sample(phishing["phishing"], settings["phishing_sample"])),
+            "reference, phishing": keep(reference["train"]["phishing"]),
+        }
         results["phishing_detection"] = train_classifier(
-            "phishing_detection", clean_urls(positives), clean_urls(negatives), LEXICAL_FEATURES, settings,
+            "phishing_detection", positives, benign, lexical, settings,
             extra_tests={PHISHING_REFERENCE + " (test split)": {"negatives": reference["test"]["legitimate"], "positives": reference["test"]["phishing"]}})
 
-    if "malware_detection" in try_models:
+    if "malware_detection" in wanted:
         progress("training malware_detection")
-        positives = sample(collection["malware"], settings["malware_sample"])
-        results["malware_detection"] = train_classifier(
-            "malware_detection", positives, keep(benign + reference["train"]["legitimate"]), LEXICAL_FEATURES, settings)
+        positives = {"older collection, malware": sample(collection["malware"], settings["malware_sample"])}
+        results["malware_detection"] = train_classifier("malware_detection", positives, benign, lexical, settings)
 
-    if "web_attack_detection" in try_models:
+    if "web_attack_detection" in wanted:
         progress("training web_attack_detection")
         if not payloads["attack"]:
             raise TrainingError("No attack payloads were loaded")
-        bases = sample(collection["benign"], settings["attack_benign_sample"])
-        with_query = [url for url in collection["benign"] if "?" in url]
-        ordinary = clean_urls(bases + sample(with_query, settings["attack_benign_sample"]) + reference["train"]["legitimate"])
-        ordinary += [compose_url(rng.choice(bases), text, rng) for text in payloads["normal"]]
+        bases = sample(all_benign, settings["attack_benign_sample"])
+        with_query = [url for url in phishing["benign"] + collection["benign"] if "?" in url]
+        ordinary = {
+            **benign,
+            "benign URLs with a query": keep(sample(with_query, settings["attack_benign_sample"])),
+            "normal strings placed like payloads": [compose_url(rng.choice(bases), text, rng) for text in payloads["normal"]],
+        }
         # each payload is carried on two different sites, so the model cannot learn a site
-        attacks = [compose_url(rng.choice(bases), text, rng) for text in payloads["attack"] for _ in range(2)]
-        results["web_attack_detection"] = train_classifier(
-            "web_attack_detection", clean_urls(attacks), clean_urls(ordinary), ATTACK_FEATURES, settings)
+        attacks = {"payloads placed on benign sites": [compose_url(rng.choice(bases), text, rng) for text in payloads["attack"] for _ in range(2)]}
+        results["web_attack_detection"] = train_classifier("web_attack_detection", attacks, ordinary, attack, settings)
         results["web_attack_detection"]["report"]["training"]["note"] = (
             "attack rows are real SQL-injection and XSS payloads placed as a query value or path segment on benign sites")
 
-    if "anomaly_detection" in try_models:
+    if "anomaly_detection" in wanted:
         progress("training anomaly_detection")
-        malicious = sample(collection["phishing"], 10000) + sample(collection["malware"], 10000)
-        results["anomaly_detection"] = train_anomaly(keep(benign), malicious, LEXICAL_FEATURES, settings)
+        malicious = sample(phishing["phishing"], 10000) + sample(collection["malware"], 10000)
+        results["anomaly_detection"] = train_anomaly(clean_urls(all_benign), malicious, lexical, settings)
 
     trained_at = datetime.utcnow().isoformat() + "Z"
     reports, accepted = {}, {}
     for name, result in results.items():
         report = result["report"]
         report.update({"name": name, "trained_at": trained_at, "settings": settings, "sources": sources,
-                       "feature_set": "url_features.py, scheme-free"})
+                       "feature_set": "url_features.py, scheme-free and without a leading www."})
         reports[name] = report
         if report["gate"]["passed"]:
             accepted[name] = result["model"]

url_features.py

@@ -110,8 +110,12 @@ def split_url(url: str) -> Dict:
 
 
 def bare_url(url: str) -> str:
-    """The URL as the models see it: no scheme, no surrounding space."""
-    return SCHEME.sub("", str(url).strip(), count=1)
+    """The URL as the models see it: no scheme, no leading "www.", no
+    surrounding space. Like the scheme, "www." says how a list was written, not
+    what the site is: one collection writes every benign URL without it and
+    another with it."""
+    bare = SCHEME.sub("", str(url).strip(), count=1)
+    return bare[4:] if bare[:4].lower() == "www." and len(bare) > 4 else bare
 
 
 def _entropy(text: str) -> float: