#!/usr/bin/env python3
"""
Check Echo's self-revealing ledger seals (`sha256-live-v5`) without trusting Echo.

    python3 check.py SEAL_DIR [SEAL_DIR ...] [--tle PATH] [--offline]

Needs Python 3.11+, the `cryptography` package (ChaCha20-Poly1305) and drand's `tle` binary
(`go install github.com/drand/tlock/cmd/tle@v1.2.0`). It imports nothing of Echo's: a checker that needs the thing it
checks is not a checker, which is why it is copied, whole, into every seal directory.

Each SEAL_DIR is one seal: `manifest.txt` (the exact text the seal's digest is the sha256 of), `seal.json`, and per
reveal group `r<ROUND>.tlock` (the group's key, time-locked to that drand round) and `r<ROUND>.ct` (the group,
encrypted under that key), and, once the group has been revealed, `r<ROUND>.held.json` (how many of its withheld
lines Echo publishes no salt for, and why) and any number of `r<ROUND>.salts.<N>.json`: each the salts of withheld
lines Echo decided to open and had not erased, each with its line. `manifest.txt.ots`, when present, is an
OpenTimestamps proof over the manifest: `ots verify manifest.txt.ots` checks it against Bitcoin.

What it checks, from the files alone and the beacon:
  1. the manifest hashes to the seal's digest, and each seal names the one before it (when several are given);
  2. each group's ciphertext and locked key hash to what the manifest commits, and the key is locked to the
     manifest's round of the manifest's chain;
  3. for every group whose round the beacon has signed: the round's signature as each relay serves it, the key
     unlocked with it, the group decrypted, its bytes hashed to the committed digest, its lines and withheld lines
     counted against the committed counts, its lines in id order, and every line's reveal day recomputed from its
     own settling instant under the seal's named rule (`utc-day-after-close`: the first 00:00 UTC after
     `settles_at`, else after the deadline);
  4. for every opened group with releases: each names this seal, round and group and its own number; each released
     salt and line hash to the withheld line's committed hash; no line is settled twice; and a hold, where there is
     one, names the group and counts no more lines than are withheld;
  5. once every group of every seal from the first is open and every withheld line released: the union of their v4
     lines, sorted by id, rehashes to the v4 digest the last seal follows, so no sealed forecast was left out of a
     reveal. Where a withheld line has no released salt (held, erased, or not yet released) that digest cannot be
     rebuilt, and this says so and why: what is checked then is only that every forecast the v4 seal covers is in a
     group, counted. A withheld line is never called verified.

A group whose round is not yet signed, or whose withheld lines have no release yet, is reported as pending or held,
not as a failure. Exit 0 when everything checkable passed, 1 when anything did not match.
"""
from __future__ import annotations

import argparse
import base64
import hashlib
import json
import re
import subprocess
import sys
import urllib.request
import zlib
from datetime import datetime, timedelta, timezone
from pathlib import Path


def sha(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()


def parse_manifest(text: str) -> dict:
    lines = text.split("\n")
    if lines[0] != "sha256-live-v5":
        raise ValueError(f"not a v5 manifest: {lines[0]!r}")
    m = {"groups": []}
    for ln in lines[1:]:
        word, _, rest = ln.partition(" ")
        if word == "prev":
            m["prev"] = None if rest == "-" else rest
        elif word == "follows":
            seq, digest, n = rest.split(" ")
            m["follows"] = {"seq": int(seq), "digest": digest, "n": int(n)}
        elif word == "covers":
            m["covers"] = rest.split(" ")
        elif word == "variant":
            m["variant"] = json.loads(rest)
        elif word == "counts":
            m["n"], m["n_withheld"], m["n_groups"] = map(int, rest.split(" "))
        elif word == "group":
            r, at, n, w, digest, ct, key, late = rest.split(" ")
            m["groups"].append({"round": int(r), "reveal_at": at, "n": int(n), "n_withheld": int(w),
                                "digest": digest, "ct": ct, "key": key, "late": late == "1"})
        else:
            raise ValueError(f"unknown manifest line {ln!r}")
    return m


def stanza(locked: str) -> tuple[int, str] | None:
    body = "".join(ln for ln in locked.strip().splitlines() if not ln.startswith("-----"))
    m = re.search(rb"-> tlock (\d+) ([0-9a-f]{64})", base64.b64decode(body))
    return (int(m.group(1)), m.group(2).decode()) if m else None


def round_at(beacon: dict, at: datetime) -> int:
    t = at.timestamp()
    whole = int(t) + (t != int(t))
    if whole <= beacon["genesis_time"]:
        return 1
    return -(-(whole - beacon["genesis_time"]) // beacon["period"]) + 1


def reveal_day(closes: datetime) -> datetime:
    d = closes.astimezone(timezone.utc).date() + timedelta(days=1)
    return datetime(d.year, d.month, d.day, tzinfo=timezone.utc)


def parse_line(ln: str) -> dict:
    """One line: `id|issued_at|deadline|forecaster|p|settlement|anchor<TAB>open:<v4>` or `...<TAB>withheld:<hex>`."""
    numeric, tab, text = ln.partition("\t")
    if not tab:
        raise ValueError("a line without its text layer")
    fid, issued, deadline, forecaster, p, rest = numeric.split("|", 5)
    dec = json.JSONDecoder()
    settle, end = dec.raw_decode(rest)
    if rest[end] != "|":
        raise ValueError("settlement not followed by the anchor")
    anchor, end2 = dec.raw_decode(rest, end + 1)
    if end2 != len(rest):
        raise ValueError("trailing bytes after the anchor")
    out = {"id": int(fid), "issued_at": issued, "deadline": deadline, "forecaster": forecaster, "p": float(p),
           "settlement": settle, "anchor": anchor}
    if text.startswith("open:"):
        out["v4"] = text[5:]
        if not out["v4"].startswith(f"{fid}|"):
            raise ValueError("the text layer is another forecast's line")
    elif re.fullmatch(r"withheld:[0-9a-f]{64}", text):
        out["withheld"] = text[9:]
    else:
        raise ValueError(f"unknown text layer {text[:20]!r}")
    return out


def beacon_round(beacon: dict, rnd: int) -> dict:
    """What each relay serves for the round: its signature, or why there is none."""
    out = {}
    for relay in beacon["relays"]:
        url = f"{relay.rstrip('/')}/{beacon['chain_hash']}/public/{rnd}"
        try:
            with urllib.request.urlopen(url, timeout=20) as r:                 # noqa: S310 — the seal's own relays
                out[relay] = json.loads(r.read()).get("signature")
        except Exception as exc:                                               # noqa: BLE001 — report, go on
            out[relay] = f"unavailable ({type(exc).__name__})"
    return out


def unlock(tle: str, beacon: dict, locked: str) -> bytes | None:
    """The key, or None while the round is unsigned. Each relay is tried in turn."""
    errors = []
    for relay in beacon["relays"]:
        r = subprocess.run([tle, "-d", "-n", relay.rstrip("/") + "/", "-c", beacon["chain_hash"]],
                           input=locked.encode(), capture_output=True, timeout=120, check=False)
        if r.returncode == 0:
            return r.stdout
        if b"too early" in r.stderr:
            return None
        errors.append(f"{relay}: {r.stderr.decode(errors='replace').strip()[-160:]}")
    raise RuntimeError("; ".join(errors))


def release(d: Path, seal_digest: str, g: dict, lines: list[dict], say) -> tuple[dict[int, str], int, str | None]:
    """An opened group's releases and hold, checked against the group: (id -> the v4 line each released salt opens,
    how many withheld lines were opened and found erased, and the hold's reason where Echo published one)."""
    r = g["round"]
    shut = {x["id"]: x["withheld"] for x in lines if "withheld" in x}
    out, erased, seen, bad = {}, set(), set(), []
    files = sorted(d.glob(f"r{r}.salts.*.json"), key=lambda p: int(p.name.split(".")[-2]))
    for k, path in enumerate(files):
        rel = json.loads(path.read_text())
        say(rel.get("format") == "echo-reveal-salts/1" and rel.get("seal") == seal_digest
            and rel.get("round") == r and rel.get("group") == g["digest"] and rel.get("n") == k
            and path.name == f"r{r}.salts.{k}.json",
            f"round {r}: release {k} names this seal, round and group, and follows the one before")
        for x in rel.get("released", []):
            i, line = x["id"], x["line"]
            opens = (i in shut and i not in seen and line.startswith(f"{i}|")
                     and sha(bytes.fromhex(x["salt"]) + line.encode()) == shut[i])
            if opens:
                out[i] = line
            else:
                bad.append(i)
            seen.add(i)
        for i in rel.get("erased", []):
            if i not in shut or i in seen:
                bad.append(i)
            erased.add(i)
            seen.add(i)
    if files:
        say(not bad, f"round {r}: {len(out)} released lines open their committed hashes, {len(erased)} opened lines "
                     f"found erased, none settled twice" + (f" ({len(bad)} do not, e.g. {bad[:3]})" if bad else ""))
    why = None
    hold = d / f"r{r}.held.json"
    if hold.exists():
        h = json.loads(hold.read_text())
        say(h.get("format") == "echo-reveal-held/1" and h.get("seal") == seal_digest and h.get("round") == r
            and h.get("group") == g["digest"] and 0 <= h.get("held", -1) <= len(shut),
            f"round {r}: the hold names this seal, round and group, and counts at most its withheld lines")
        why = str(h.get("why") or "no reason given")
    left = len(shut) - len(out) - len(erased)
    if left:
        print(f"  ....  round {r}: {left} withheld lines have no released salt: counted, not rehashed"
              + (f" (Echo publishes none: {why})" if why else " (no release or hold published yet)"))
    return out, len(erased), why


def check(dirs: list[Path], tle: str = "tle", offline: bool = False) -> int:
    from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305

    failed, pending, withheld, held, erased, union = 0, 0, 0, 0, 0, {}
    reasons: set[str] = set()
    first_prev = "unset"

    def say(ok: bool, what: str) -> None:
        nonlocal failed
        failed += not ok
        print(f"  {'PASS' if ok else 'FAIL'}  {what}")

    prev_digest = None
    seals = []
    for d in dirs:
        text = (d / "manifest.txt").read_bytes()
        row = json.loads((d / "seal.json").read_text())
        seals.append((d, text, row, parse_manifest(text.decode())))
    seals.sort(key=lambda s: s[2]["seq"])
    last = None
    for i, (d, text, row, m) in enumerate(seals):
        print(f"v5 seal {row['seq']}: {m['n']} forecasts in {m['n_groups']} reveal groups, {m['n_withheld']} withheld, "
              f"follows v4 seal {m['follows']['seq']}")
        say(sha(text) == row["digest"], f"sha256(manifest.txt) is the seal's digest {row['digest'][:16]}…")
        say(len(m["groups"]) == m["n_groups"] and sum(g["n"] for g in m["groups"]) == m["n"]
            and sum(g["n_withheld"] for g in m["groups"]) == m["n_withheld"], "the manifest's counts add up")
        if i == 0:
            first_prev = m["prev"]
        else:
            say(m["prev"] == prev_digest, "it names the seal before it")
        prev_digest, last = row["digest"], m
        b = m["variant"]["beacon"]
        for g in m["groups"]:
            r = g["round"]
            ct, locked = (d / f"r{r}.ct").read_bytes(), (d / f"r{r}.tlock").read_text()
            say(sha(ct) == g["ct"] and sha(locked.encode()) == g["key"],
                f"round {r}: the ciphertext and the locked key are the committed ones")
            say(stanza(locked) == (r, b["chain_hash"]), f"round {r}: the key is locked to that round of that chain")
            say(round_at(b, datetime.fromisoformat(g["reveal_at"])) == r, f"round {r} is signed at {g['reveal_at']}")
            if not offline:
                sigs = beacon_round(b, r)
                good = {s for s in sigs.values() if s and not s.startswith("unavailable")}
                print(f"  ....  round {r} signature from {len(sigs)} relays: "
                      + ("not yet signed" if not good else f"{len(good)} distinct: {min(good)[:24]}…"))
                if len(good) > 1:
                    say(False, f"round {r}: relays disagree on its signature")
            key = unlock(tle, b, locked)
            if key is None:
                pending += 1
                print(f"  ....  round {r} ({g['reveal_at']}) not yet signed: {g['n']} forecasts stay locked")
                continue
            try:
                plain = zlib.decompress(ChaCha20Poly1305(key).decrypt(b"\x00" * 12, ct, None))
            except Exception as exc:                                       # noqa: BLE001 — a finding, not a crash
                say(False, f"round {r}: the ciphertext does not open under its key ({type(exc).__name__})")
                continue
            say(sha(plain) == g["digest"], f"round {r}: the opened group hashes to its committed digest")
            lines = [parse_line(ln) for ln in plain.decode().split("\n")]
            ids = [x["id"] for x in lines]
            say(len(lines) == g["n"] and sum("withheld" in x for x in lines) == g["n_withheld"],
                f"round {r}: {len(lines)} lines, {g['n_withheld']} withheld, as committed")
            say(ids == sorted(set(ids)), f"round {r}: lines in id order, none twice")
            due = datetime.fromisoformat(g["reveal_at"])
            if m["variant"].get("reveal") == "utc-day-after-close":
                wrong = [x["id"] for x in lines if reveal_day(datetime.fromisoformat(
                    (x["settlement"] or {}).get("settles_at") or x["deadline"])) != due]
                say(not wrong, f"round {r}: every line belongs to this reveal day"
                    + (f" ({len(wrong)} do not, e.g. {wrong[:3]})" if wrong else ""))
            else:
                print(f"  ....  round {r}: reveal rule {m['variant'].get('reveal')!r} is not one this checker "
                      "recomputes; which lines belong to the day is unchecked")
            opened, gone, why = release(d, row["digest"], g, lines, say)
            for x in lines:
                if x["id"] in union:
                    say(False, f"forecast {x['id']} is in two groups")
                union[x["id"]] = x.get("v4") or opened.get(x["id"])
            shut = sum(1 for x in lines if "withheld" in x and x["id"] not in opened) - gone
            erased += gone
            if why is None:
                withheld += shut
            else:
                held += shut
                reasons.add(why)
    if last is not None:
        if first_prev is not None:
            print("  ....  completeness: the first seal given is not the chain's first; give every seal from it")
        elif pending:
            print(f"  ....  completeness: {pending} groups still locked; it is checkable once all have opened")
        elif withheld or held or erased:
            told = [f"{withheld} withheld lines not yet released"] if withheld else []
            told += ([f"{held} withheld lines held, for which Echo publishes no salt ({'; '.join(sorted(reasons))})"]
                     if held else [])
            told += [f"{erased} erased stay hashes"] if erased else []
            print(f"  ....  completeness: {', '.join(told)}: those lines can be counted but not rehashed, so the v4 "
                  f"lines cannot all be rebuilt and the v4 digest is NOT checked; {len(union)} forecasts against "
                  f"the v4 seal's {last['follows']['n']}")
            say(len(union) == last["follows"]["n"], "every forecast the v4 seal covers is in a group, counted")
        else:
            got = sha("\n".join(union[k] for k in sorted(union)).encode())
            say(got == last["follows"]["digest"] and len(union) == last["follows"]["n"],
                f"the union of {len(union)} lines rehashes to v4 seal {last['follows']['seq']}'s digest")
    print("\nNOT VERIFIED" if failed else f"\nVERIFIED{f' ({pending} groups not yet revealed)' if pending else ''}")
    return 1 if failed else 0


def main(argv: list[str]) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("dirs", nargs="*", default=[str(Path(__file__).resolve().parent)],
                    help="seal directories, any order (default: the one this script sits in)")
    ap.add_argument("--tle", default="tle", help="the tle binary that unlocks a key once its round is signed")
    ap.add_argument("--offline", action="store_true",
                    help="do not fetch each round's signature from the relays for display; tle still asks them")
    args = ap.parse_args(argv)
    return check([Path(d) for d in args.dirs], tle=args.tle, offline=args.offline)


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))