#!/usr/bin/env python3
"""How often is granted access ever removed? Measured on one public provisioning log, by year of first grant.

The log is the UCI Machine Learning Repository dataset 216, "Amazon Access Samples" (CC BY 4.0, donated
2011-09-12): 716,063 events in which a person (LOGIN, an anonymized number) was given or lost access
(ACTION) to a resource (TARGET_NAME, an anonymized number), from 2005-03-25 to 2010-08-31.
https://archive.ics.uci.edu/dataset/216/amazon+access+samples

    python3 access_removal_by_year.py --data-dir <empty folder> --out access_removal_by_year.csv
    python3 access_removal_by_year.py --data-dir <another empty folder> --out planted.csv \
        --plant duplicate-block,remove-before-add,unknown-action,bad-request-date,old-date

Needs Python 3.12 or later (or 3.11.4, 3.10.12, 3.9.17 and up, which have tarfile's extraction filter)
and the duckdb module (tested on Python 3.12, DuckDB 1.5.6). It downloads one 12 MB zip,
checks its byte size and SHA-256 (the server sends no Last-Modified header, so these are the pin), opens
the .tgz inside it with tarfile and extracts only the history file. The other member of the .tgz is 4.8 GB
and is never read. It prints counts and shares only: no login or resource number is printed or written.
It exits 1, and writes no CSV, if a check fails.

The unit is a pair: one LOGIN and one TARGET_NAME. A pair is "granted" from its first add_access event.
It is "removed later" if any remove_access event for the same pair comes after that first grant. The time
of an event is REQUEST_DATE, which parses in every row. (AUTHORIZATION_DATE does not: see the printed count.)
"""
import argparse
import csv
import hashlib
import os
import sys
import tarfile
import urllib.request
import zipfile

import duckdb

URL = "https://archive.ics.uci.edu/static/public/216/amazon+access+samples.zip"
ZIP_BYTES = 12_268_663
ZIP_SHA256 = "d493e3378508b7f3141253c421c2267923fd8ca60ea4088afe72b21b8228d224"
TGZ = "amzn-anon-access-samples.tgz"
HISTORY = "amzn-anon-access-samples-history-2.0.csv"
HISTORY_BYTES = 44_583_050

# What each ACTION value means. An action that is not listed here is counted as unmapped and fails the run.
ACTIONS = {"add_access": "grant", "remove_access": "removal"}

# Years of first grant that the log covers (the first event is on 2005-03-25, the last on 2010-08-31).
YEARS = [2005, 2006, 2007, 2008, 2009, 2010]
WINDOWS_DAYS = [30, 90, 365]   # "removed within N days"; judged only on pairs followed for the full 365 days

# Facts about the published file, from our first read of it. A copy that differs from them is not the file
# this page measured, or has been damaged.
PINNED = {"rows": 716_063, "add_access": 705_152, "remove_access": 10_911, "repeated_keys": 6_283,
          "removal_before_first_grant": 40, "first_event": "2005-03-25 10:53:40", "last_event": "2010-08-31 18:03:18"}

PLANTS = {"duplicate-block", "remove-before-add", "unknown-action", "bad-request-date", "old-date", "no-removals"}

FIELDS = ["first_grant_year", "pairs_granted", "removed_later", "pct_removed_later", "open_at_log_end",
          "pct_open_at_log_end", "median_days_to_removal", "pairs_followed_365d"] + \
         [f"pct_removed_within_{n}d" for n in WINDOWS_DAYS]


def fetch(data_dir):
    """Download the zip, check size and SHA-256, extract only the history file. Return its path."""
    zpath = os.path.join(data_dir, "amazon-access-samples.zip")
    with urllib.request.urlopen(URL, timeout=120) as resp:
        modified = resp.headers.get("Last-Modified")
        body = resp.read()
    with open(zpath, "wb") as f:
        f.write(body)
    sha = hashlib.sha256(body).hexdigest()
    print(f"source  {URL}")
    print(f"        bytes={len(body)} sha256={sha[:16]}... last-modified={modified or 'not sent'}")
    if len(body) != ZIP_BYTES or sha != ZIP_SHA256:
        sys.exit(f"the zip is not the file this script was written for (expected {ZIP_BYTES} bytes, sha256 {ZIP_SHA256})")
    with zipfile.ZipFile(zpath) as z:
        if z.namelist() != [TGZ]:
            sys.exit(f"unexpected zip contents: {z.namelist()}")
        with z.open(TGZ) as raw, tarfile.open(fileobj=raw, mode="r|gz") as tar:
            for member in tar:
                if member.name == HISTORY:
                    tar.extract(member, os.path.join(data_dir, "x"), filter="data")
                    break
            else:
                sys.exit(f"{HISTORY} not found in {TGZ}")
    path = os.path.join(data_dir, "x", HISTORY)
    print(f"        extracted {HISTORY}: {os.path.getsize(path)} bytes (the other member of the .tgz was not read)")
    if os.path.getsize(path) != HISTORY_BYTES:
        sys.exit(f"{HISTORY} is {os.path.getsize(path)} bytes, expected {HISTORY_BYTES}")
    return path


def load(path, plant):
    with open(path, "rb") as f:
        source_lines = sum(chunk.count(b"\n") for chunk in iter(lambda: f.read(1 << 20), b"")) - 1  # minus header
    sql_path = path.replace("'", "''")  # a quote in the folder name must not end the SQL string
    con = duckdb.connect()
    con.sql(f"CREATE TABLE ev AS SELECT row_number() OVER () AS file_row, ACTION, TARGET_NAME, LOGIN, REQUEST_DATE, "
            f"AUTHORIZATION_DATE FROM read_csv('{sql_path}', header=true, all_varchar=true)")
    if "duplicate-block" in plant:
        con.sql("INSERT INTO ev SELECT * FROM ev ORDER BY file_row LIMIT 1000")
        print("PLANTED: the first 1,000 events loaded twice")
    if "remove-before-add" in plant:
        con.sql("""INSERT INTO ev
                   SELECT 800000000 + row_number() OVER (), 'remove_access', target_name, login,
                          strftime(first_add - INTERVAL 1 HOUR, '%Y-%m-%d %H:%M:%S'), NULL
                   FROM (SELECT login, target_name, MIN(TRY_STRPTIME(request_date, '%Y-%m-%d %H:%M:%S')) AS first_add
                         FROM ev GROUP BY login, target_name
                         HAVING COUNT(*) FILTER (WHERE action = 'remove_access') = 0
                         ORDER BY login, target_name LIMIT 200)""")
        print("PLANTED: 200 pairs that were only ever granted given a removal one hour before their first grant")
    if "unknown-action" in plant:
        con.sql("UPDATE ev SET action = 'revoke_access' WHERE file_row IN "
                "(SELECT file_row FROM ev WHERE action = 'remove_access' ORDER BY file_row LIMIT 300)")
        print("PLANTED: 300 remove_access events renamed 'revoke_access'")
    if "bad-request-date" in plant:
        con.sql("UPDATE ev SET request_date = 'n/a' WHERE file_row IN (SELECT file_row FROM ev "
                "WHERE file_row <= 716063 ORDER BY file_row DESC LIMIT 100)")
        print("PLANTED: the request date of the last 100 events of the file replaced by 'n/a'")
    if "old-date" in plant:
        con.sql("UPDATE ev SET request_date = '1999-01-01 00:00:00' WHERE file_row IN "
                "(SELECT file_row FROM ev WHERE action = 'add_access' ORDER BY file_row LIMIT 50)")
        print("PLANTED: 50 grants dated 1999-01-01")
    if "no-removals" in plant:
        con.sql("DELETE FROM ev WHERE action = 'remove_access'")
        print("PLANTED: every remove_access event deleted")
    return con, source_lines


def measure(con):
    """Build pairs and cohorts in DuckDB. Return (result rows, facts about the log)."""
    con.sql("CREATE TABLE d AS SELECT file_row, action, login, target_name, "
            "TRY_STRPTIME(request_date, '%Y-%m-%d %H:%M:%S') AS ts FROM ev")
    first_ts, last_ts = con.sql("SELECT MIN(ts), MAX(ts) FROM d").fetchone()
    # one row per pair; first grant, and the first removal that comes after it
    con.sql("""CREATE TABLE pairs AS
               SELECT login, target_name,
                      MIN(ts) FILTER (WHERE action = 'add_access') AS first_add,
                      MIN(ts) FILTER (WHERE action = 'remove_access') AS first_remove_any
               FROM d GROUP BY login, target_name""")
    con.sql("""CREATE TABLE removed AS
               SELECT p.login, p.target_name, MIN(d.ts) AS first_removal
               FROM pairs p JOIN d ON d.login = p.login AND d.target_name = p.target_name
                                  AND d.action = 'remove_access' AND d.ts > p.first_add
               GROUP BY p.login, p.target_name""")
    # the last event of each pair decides whether access is still open at the end of the log
    con.sql("""CREATE TABLE last_ev AS
               SELECT login, target_name, action AS last_action FROM (
                 SELECT *, row_number() OVER (PARTITION BY login, target_name ORDER BY ts DESC, file_row DESC) AS rn
                 FROM d) WHERE rn = 1""")
    con.sql("""CREATE TABLE cohort AS
               SELECT YEAR(p.first_add) AS yr, p.first_add, r.first_removal, l.last_action
               FROM pairs p LEFT JOIN removed r USING (login, target_name) JOIN last_ev l USING (login, target_name)
               WHERE p.first_add IS NOT NULL""")
    ends = f"TIMESTAMP '{last_ts}'"
    followed = f"first_add <= {ends} - INTERVAL 365 DAY"
    within = ", ".join(
        f"COUNT(*) FILTER (WHERE {followed} AND first_removal <= first_add + INTERVAL {n} DAY) AS w{n}"
        for n in WINDOWS_DAYS)
    rows = con.sql(f"""
        SELECT yr, COUNT(*) AS pairs, COUNT(first_removal) AS removed,
               COUNT(*) FILTER (WHERE last_action = 'add_access') AS open_end,
               quantile_cont(date_diff('second', first_add, first_removal) / 86400.0, 0.5) AS median_days,
               COUNT(*) FILTER (WHERE {followed}) AS followed, {within}
        FROM cohort GROUP BY ROLLUP (yr) ORDER BY yr NULLS LAST""").fetchall()
    facts = {"first_event": first_ts.strftime("%Y-%m-%d %H:%M:%S"), "last_event": last_ts.strftime("%Y-%m-%d %H:%M:%S")}
    q = lambda sql: con.sql(sql).fetchone()[0]
    facts["rows"] = q("SELECT COUNT(*) FROM ev")
    facts["add_access"] = q("SELECT COUNT(*) FROM ev WHERE action = 'add_access'")
    facts["remove_access"] = q("SELECT COUNT(*) FROM ev WHERE action = 'remove_access'")
    facts["bad_request_dates"] = q("SELECT COUNT(*) FROM d WHERE ts IS NULL")
    facts["bad_authorization_dates"] = q("SELECT COUNT(*) FROM ev WHERE TRY_STRPTIME(authorization_date, "
                                         "'%Y-%m-%d %H:%M:%S') IS NULL")
    facts["repeated_keys"] = q("SELECT COUNT(*) FROM (SELECT 1 FROM ev GROUP BY action, target_name, login, "
                               "request_date HAVING COUNT(*) > 1)")
    facts["pairs_all"] = q("SELECT COUNT(*) FROM pairs")
    facts["pairs_granted"] = q("SELECT COUNT(*) FROM pairs WHERE first_add IS NOT NULL")
    facts["removal_only_pairs"] = q("SELECT COUNT(*) FROM pairs WHERE first_add IS NULL")
    facts["removal_before_first_grant"] = q("SELECT COUNT(*) FROM pairs WHERE first_remove_any < first_add")
    facts["logins"] = q("SELECT COUNT(DISTINCT login) FROM ev")
    facts["resources"] = q("SELECT COUNT(DISTINCT target_name) FROM ev")
    # an independent count of pairs removed after their first grant, by a different query (EXISTS, no join table)
    facts["removed_independent"] = q("""SELECT COUNT(*) FROM pairs p WHERE p.first_add IS NOT NULL AND EXISTS (
                                          SELECT 1 FROM d WHERE d.login = p.login AND d.target_name = p.target_name
                                            AND d.action = 'remove_access' AND d.ts > p.first_add)""")
    return rows, facts


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-dir", required=True, help="a new, empty folder for the download")
    ap.add_argument("--out", required=True, help="CSV to write")
    ap.add_argument("--plant", default="", help=f"comma list of test faults: {', '.join(sorted(PLANTS))}")
    a = ap.parse_args()
    plant = [x for x in a.plant.split(",") if x]
    if set(plant) - PLANTS:
        ap.error(f"unknown --plant value(s) {sorted(set(plant) - PLANTS)}; use {sorted(PLANTS)}")
    os.makedirs(a.data_dir, exist_ok=True)
    if os.listdir(a.data_dir):
        sys.exit(f"{a.data_dir} is not empty: give the script a new, empty folder")
    print(f"duckdb {duckdb.__version__}, python {sys.version.split()[0]}")

    path = fetch(a.data_dir)
    con, source_lines = load(path, plant)
    raw, facts = measure(con)
    unmapped = dict(con.sql(f"SELECT action, COUNT(*) FROM ev WHERE action NOT IN "
                            f"({', '.join(repr(k) for k in ACTIONS)}) GROUP BY 1 ORDER BY 1").fetchall())
    keys_with_rows = {k for (k,) in con.sql("SELECT DISTINCT action FROM ev").fetchall()}

    result = []
    for yr, pairs, removed, open_end, median, followed, *w in raw:
        rec = {"first_grant_year": yr if yr is not None else "all", "pairs_granted": pairs, "removed_later": removed,
               "pct_removed_later": 100.0 * removed / pairs, "open_at_log_end": open_end,
               "pct_open_at_log_end": 100.0 * open_end / pairs,
               "median_days_to_removal": median, "pairs_followed_365d": followed}
        for n, c in zip(WINDOWS_DAYS, w):
            rec[f"pct_removed_within_{n}d"] = 100.0 * c / followed if followed else None
        result.append(rec)

    print(f"\nevents loaded: {facts['rows']}  (add_access {facts['add_access']}, remove_access {facts['remove_access']})")
    print(f"first event {facts['first_event']}, last event {facts['last_event']}")
    print(f"logins {facts['logins']}, resources {facts['resources']}, login-resource pairs {facts['pairs_all']}")
    print(f"pairs ever granted {facts['pairs_granted']}; pairs with a removal and no grant in the log "
          f"{facts['removal_only_pairs']}")
    print(f"pairs with a removal before their first grant {facts['removal_before_first_grant']}")
    print(f"events with an AUTHORIZATION_DATE that is not a timestamp {facts['bad_authorization_dates']}; "
          f"with a REQUEST_DATE that is not {facts['bad_request_dates']}")
    print(f"action-resource-login-time keys that occur more than once {facts['repeated_keys']}")

    def fmt(v, nd=2):
        return "n/a" if v is None else f"{v:.{nd}f}"
    print()
    head = ["year", "pairs", "removed", "%removed", "open_end", "%open", "median_d", "followed", "%in30d", "%in90d",
            "%in365d"]
    print("  ".join(h.rjust(9) for h in head))
    for r in result:
        cells = [str(r["first_grant_year"]), str(r["pairs_granted"]), str(r["removed_later"]),
                 fmt(r["pct_removed_later"]), str(r["open_at_log_end"]), fmt(r["pct_open_at_log_end"]),
                 fmt(r["median_days_to_removal"], 1), str(r["pairs_followed_365d"])] + \
                [fmt(r[f"pct_removed_within_{n}d"]) for n in WINDOWS_DAYS]
        print("  ".join(c.rjust(9) for c in cells))
    print()

    # The script's own checks. Each names the planted fault that trips it.
    failures = []
    if facts["rows"] != source_lines:   # duplicate-block, no-removals
        failures.append(f"loaded {facts['rows']} events but the file has {source_lines} data lines "
                        f"({facts['rows'] - source_lines:+d})")
    for k in ("rows", "add_access", "remove_access", "repeated_keys", "removal_before_first_grant"):
        if facts[k] != PINNED[k]:   # duplicate-block, no-removals, unknown-action, remove-before-add
            failures.append(f"{k} is {facts[k]}, the published file has {PINNED[k]} ({facts[k] - PINNED[k]:+d})")
    for k in ("first_event", "last_event"):
        if facts[k] != PINNED[k]:   # old-date, bad-request-date
            failures.append(f"{k} is {facts[k]}, the published file has {PINNED[k]}")
    if unmapped:   # unknown-action
        failures.append("events with an action that is not in ACTIONS: " +
                        ", ".join(f"'{k}' x {v}" for k, v in unmapped.items()))
    for k in ACTIONS:   # no-removals
        if k not in keys_with_rows:
            failures.append(f"ACTIONS key '{k}' matches no event in the file")
    if facts["bad_request_dates"]:   # bad-request-date
        failures.append(f"{facts['bad_request_dates']} events have a REQUEST_DATE that is not a timestamp")
    by_year = {r["first_grant_year"]: r for r in result if r["first_grant_year"] != "all"}
    allrow = next(r for r in result if r["first_grant_year"] == "all")
    outside = {y: r["pairs_granted"] for y, r in by_year.items() if y not in YEARS}
    if outside:   # old-date
        failures.append(f"pairs first granted in a year that is not in YEARS: {outside}")
    for y in YEARS:   # a year in YEARS that no pair starts in would mean a wrong list
        if y not in by_year:
            failures.append(f"YEARS lists {y} but no pair was first granted then")
    # a join that drops or doubles pairs would make the two counts differ (no run plants this fault)
    if allrow["removed_later"] != facts["removed_independent"]:
        failures.append(f"removed pairs {allrow['removed_later']} differ from a count by a second query "
                        f"{facts['removed_independent']}")

    if failures:
        print("SELF-CHECK FAILED: no CSV written")
        for f in failures:
            print("  -", f)
        sys.exit(1)
    print(f"SELF-CHECK PASSED: {facts['rows']} events equal the {source_lines} data lines of the file and the published "
          f"counts, every action is mapped, every ACTIONS key matches events, every REQUEST_DATE parses, "
          f"every first grant falls in {YEARS[0]}-{YEARS[-1]}, removed pairs equal a count by a second query")
    with open(a.out, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=FIELDS, lineterminator="\n")
        w.writeheader()
        for r in result:
            w.writerow({k: ("" if r[k] is None else f"{r[k]:.2f}" if isinstance(r[k], float) else r[k]) for k in FIELDS})
    print(f"wrote {a.out}")


if __name__ == "__main__":
    main()