Security alert triage

"""Security alert triage: which earlier event was the entry point of an account takeover?

Simulates 1,500 employee accounts over 30 days of identity telemetry: logins (new device,
new country, hosting ASN, MFA method, MFA push fatigue), OAuth consent grants, password
resets, phishing clicks reported by the mail gateway, new inbox rules and token use from a
new IP. Some accounts are taken over; the takeover is confirmed minutes to days later.

Hidden ground truth (the seeded entry vectors, per event):
  * oauth:docs-viewer-pro  -- a consent-phishing app; a grant leads to takeover with p=0.55
  * reset:helpdesk         -- a social-engineered helpdesk reset; p=0.20 (most are legit)
  * plus a small background rate of takeovers with no visible entry event
Distractors: travel logins from new countries (common, benign), new devices after phone
upgrades (common, benign). Honest near-miss: the consent-phishing email is often clicked
(phish_click) minutes before the malicious grant, but the click itself does nothing here.
After a takeover the attacker logs in from a hosting ASN, uses tokens from a new IP and adds
a forwarding inbox rule: symptoms that precede the confirmation but are not entry points.
Confounder: the malicious app always asks for scope:mail.read, so that scope picks up
credit and legitimate mail apps that request it inherit some risk.

Learning: analyst verdicts on 70% of incidents -> feedback (confirm the entry event, reject
the other suspects); the rest are confirmed without a verdict -> observe; accounts with no
labelled incident -> settle once their windows have passed.

Run:
    pip install heapdb
    python 0009-security-alert-triage.py
"""

from __future__ import annotations

import math
import random

from heapdb import HeapDB, LogNormalKernel, format_result

M, H, D = 60.0, 3600.0, 86400.0
SEED = 9
N_ACCOUNTS = 1500
DAYS = 30
EFFECT = "account_takeover"
FAKE_APP = "oauth:docs-viewer-pro"
SEEDED = {FAKE_APP: 0.55, "reset:helpdesk": 0.20}
BACKGROUND = 0.03           # chance an account is taken over with no visible entry event
LABELLED = 0.70             # share of incidents that get an analyst verdict
KERNEL = LogNormalKernel(median=6 * H, sigma=1.0, tail_sigmas=3.0)  # support ~5 days
APPS = ["oauth:calendar-sync", "oauth:notes-app", "oauth:mail-merge", "oauth:crm-connector"]


def lag(rng: random.Random) -> float:
    while True:  # LogNormal lag, median 6h, kept inside the kernel's support
        x = math.exp(rng.gauss(math.log(6 * H), 1.0))
        if 2 * M < x < KERNEL.support:
            return x


def login(rng: random.Random, travel: bool, upgrade_day: float, t: float) -> set[str]:
    f = {"login", rng.choice(["mfa:push"] * 6 + ["mfa:totp"] * 3 + ["mfa:none"])}
    if travel and rng.random() < 0.5:
        f.add("geo:new_country")
    if abs(t - upgrade_day) < D and rng.random() < 0.6:
        f.add("login:new_device")
    if rng.random() < 0.01:
        f.add("asn:hosting")   # VPNs and cloud desktops
    if f & {"mfa:push"} and rng.random() < 0.01:
        f.add("push_fatigue")
    return f


def simulate(rng: random.Random):
    """Return (events, incidents). events: (acct, t, type, features); incidents: dicts."""
    events, incidents = [], []
    for a in range(N_ACCOUNTS):
        acct = f"acct-{a:04d}"
        travel_start = rng.uniform(0, DAYS * D) if rng.random() < 0.25 else -1e9
        upgrade = rng.uniform(0, DAYS * D) if rng.random() < 0.15 else -1e9
        evs = []
        for _ in range(rng.randint(15, 40)):
            t = rng.uniform(0, DAYS * D)
            evs.append((t, "login", login(rng, abs(t - travel_start - 2 * D) < 2 * D, upgrade, t)))
        if rng.random() < 0.3:
            evs.append((rng.uniform(0, DAYS * D), "oauth_grant",
                        {rng.choice(APPS)} | ({"scope:mail.read"} if rng.random() < 0.4 else set())))
        if rng.random() < 0.06:   # consent-phishing campaign reaches this user
            t = rng.uniform(0, DAYS * D)
            evs.append((t - rng.uniform(1, 10) * M, "mail", {"phish_click"}))
            evs.append((t, "oauth_grant", {FAKE_APP, "scope:mail.read"}))
        elif rng.random() < 0.05:  # ordinary phishing-test or spam click
            evs.append((rng.uniform(0, DAYS * D), "mail", {"phish_click"}))
        if rng.random() < 0.15:
            evs.append((rng.uniform(0, DAYS * D), "password_reset", {"reset:email"}))
        if rng.random() < 0.06:
            evs.append((rng.uniform(0, DAYS * D), "password_reset", {"reset:helpdesk"}))
        evs.sort(key=lambda e: e[0])
        # Who was taken over, and through which event (first one wins).
        entry, t_entry = None, math.inf
        for i, (t, _typ, feats) in enumerate(evs):
            p = max((SEEDED.get(f, 0.0) for f in feats), default=0.0)
            if p and rng.random() < p and t < t_entry:
                entry, t_entry = i, t
        if entry is None and rng.random() < BACKGROUND:
            t_entry = rng.uniform(0, DAYS * D)
        if math.isfinite(t_entry):
            confirm = t_entry + lag(rng)
            gap = confirm - t_entry
            evs.append((t_entry + 0.2 * gap, "login", {"login", "asn:hosting", "geo:new_country",
                                                      "mfa:none"}))
            evs.append((t_entry + 0.4 * gap, "token_use", {"token:new_ip"}))
            evs.append((t_entry + 0.6 * gap, "inbox_rule", {"rule:forward_external"}))
            evs.sort(key=lambda e: e[0])
            incidents.append({"acct": acct, "at": confirm, "entry": entry,
                              "entry_t": t_entry, "labelled": rng.random() < LABELLED})
        if rng.random() < 0.03:   # benign forwarding rule (personal archive)
            evs.append((rng.uniform(0, DAYS * D), "inbox_rule", {"rule:forward_external"}))
        events += [(acct, t, typ, feats) for t, typ, feats in evs]
    incidents.sort(key=lambda x: x["at"])
    return events, incidents


def main() -> None:
    rng = random.Random(SEED)
    events, incidents = simulate(rng)
    db = HeapDB(bucket_seconds=12 * H)
    db.register_effect(EFFECT, KERNEL, default_prior=(1, 19), leak=0.05)
    # Analyst policy: forwarding rules and token reuse are what attackers do AFTER entry.
    db.add_rule(EFFECT, "post-entry symptom", None, 0.3, types={"inbox_rule", "token_use"})

    for acct, t, typ, feats in events:
        db.insert(acct, t, typ, feats)
    n_fb = n_obs = 0
    labelled_accts = set()
    for inc in incidents:
        acct, at = inc["acct"], inc["at"]
        true_id = None
        if inc["entry"] is not None:
            true_id = next(e.id for e in db.events(acct)
                           if e.time == inc["entry_t"] and e.type in ("oauth_grant", "password_reset"))
        if inc["labelled"]:
            res = db.top_causes(acct, EFFECT, at, k=5)
            rejected = [c.event.id for c in res if c.event.id != true_id]
            db.feedback(acct, EFFECT, at, confirmed_event_id=true_id, rejected_ids=rejected)
            labelled_accts.add(acct)
            n_fb += 1
        else:
            db.observe(acct, EFFECT, at)
            n_obs += 1
        inc["true_id"] = true_id
    horizon = DAYS * D + KERNEL.support + D
    for acct in db.entities:
        if acct not in labelled_accts:
            db.settle(horizon, EFFECT, entity=acct)

    seeded_n = sum(1 for i in incidents if i["entry"] is not None)
    print("SCENARIO")
    print(f"  {len(events):,} identity events for {len(db.entities):,} accounts over {DAYS} days")
    print(f"  {len(incidents)} takeovers confirmed ({seeded_n} via a seeded vector, "
          f"{len(incidents) - seeded_n} background)")
    print(f"  learning: {n_fb} analyst verdicts (feedback), {n_obs} unlabelled (observe), "
          f"benign accounts settled")

    # Where does the true entry event rank now, over all seeded incidents?
    ranks_true, pick = [], None
    for inc in incidents:
        if inc["entry"] is None:
            continue
        res = db.top_causes(inc["acct"], EFFECT, inc["at"], k=5)
        ids = [c.event.id for c in res]
        tid = inc["true_id"]
        ranks_true.append(ids.index(tid) + 1 if tid in ids else 99)
        feats = {f for e in db.events(inc["acct"]) for f in e.features}
        if (pick is None and not inc["labelled"] and ids and ids[0] == tid
                and FAKE_APP in feats and "phish_click" in feats):
            pick = inc
    print()
    print("ONE CASE, EXPLAINED (after learning; this incident got no analyst verdict)")
    print(format_result(db.top_causes(pick["acct"], EFFECT, pick["at"], k=4, explain=True)))

    print()
    print("ACROSS ALL CASES (global links, top 8 by strength)")
    rep = db.link_report(EFFECT, min_support=3)
    print(f"  {'feature':<26} {'mean':>6} {'90% CI':>15} {'lift':>6} {'n':>7}")
    for r in list(rep)[:8]:
        print(f"  {r.feature:<26} {r.strength:>6.3f} [{r.lo:.3f}, {r.hi:.3f}] "
              f"{r.lift:>6.2f} {r.n:>7.1f}")

    print()
    print("RISK OF NEW EVENTS (feature_risk, global view, before any lag)")
    new = [
        ("oauth_grant", {FAKE_APP, "scope:mail.read"}, "fresh grant to docs-viewer-pro"),
        ("password_reset", {"reset:helpdesk"}, "helpdesk password reset"),
        ("mail", {"phish_click"}, "reported phishing click"),
        ("oauth_grant", {"oauth:mail-merge", "scope:mail.read"}, "grant to a known app"),
        ("login", {"login", "geo:new_country", "mfa:push"}, "login from a new country"),
        ("login", {"login", "login:new_device", "mfa:totp"}, "login from a new device"),
    ]
    for typ, feats, label in new:
        print(f"  {label:<32} {db.feature_risk(EFFECT, feats, type=typ):.4f}")

    print()
    print("CHECK")
    order = [r.feature for r in rep]
    ranks = {f: order.index(f) + 1 for f in SEEDED}
    print("  seeded vectors rank " + ", ".join(f"{f} #{r}" for f, r in ranks.items())
          + f" of {len(order)} features")
    near = order.index("phish_click") + 1
    print(f"  near-miss phish_click ranks #{near}: it precedes the malicious grant, analysts reject it")
    n = len(ranks_true)
    hit1 = sum(r == 1 for r in ranks_true)
    hit3 = sum(r <= 3 for r in ranks_true)
    print(f"  true entry event is the #1 suspect in {hit1}/{n} seeded incidents, top-3 in {hit3}/{n}")
    if sorted(ranks.values()) != [1, 2]:
        raise SystemExit(1)


if __name__ == "__main__":
    main()