#!/usr/bin/env python3
"""Independent verifier for RandNext draws ("randnext-draw-v1" and "randnext-draw-v2").

Written from the published spec (docs/verifiable-draw-design.md, section 4), not from the
TypeScript code, so the two implementations check each other. Standard library only; if the
optional `cryptography` package is installed, the Ed25519 statement signatures are checked too.

Usage:
  python3 verify-draw.py --fetch dr_7kq2m9x4tb1c [--entries entries.txt]
  python3 verify-draw.py --record record.json --entries entries.txt [--public-key HEX]
  python3 verify-draw.py --vectors tests/unit/fixtures/draw-vectors.json

--record takes the JSON from GET /api/v1/draws/{id}. --entries takes the entry list, one per
line: the file from /api/v1/draws/{id}/entries, or the organizer's own copy (it is normalized
the same way before hashing). Exit status is 0 when every check that ran passes.

What this script cannot do with the standard library: verify drand's BLS signature. In --fetch
mode it compares the beacon with drand's own relay instead; the web verifier at
https://randnext.com/draws/verify checks the BLS signature itself.
"""

import argparse
import hashlib
import hmac
import json
import struct
import sys
import unicodedata
import urllib.request

V1 = "randnext-draw-v1"
V2 = "randnext-draw-v2"
API_BASE = "https://randnext.com/api/v1/draws/"
DRAND_QUICKNET = "https://api.drand.sh/52db9ba70e0cc0f6eaf7803dd07447a1f5477735fd3f661792ba94600c84e971/public/"

# The exact whitespace set trimmed from each line (spec 4.1).
TRIM_CHARS = (
    "\t\n\v\f\r   "
    + "".join(chr(c) for c in range(0x2000, 0x200B))
    + "

  　"
)


def normalize_entries(text, duplicates):
    lines = []
    for line in text.split("\n"):
        line = unicodedata.normalize("NFC", line).strip(TRIM_CHARS)
        if line:
            lines.append(line)
    if duplicates == "merge":
        seen = set()
        merged = []
        for line in lines:
            if line not in seen:
                seen.add(line)
                merged.append(line)
        lines = merged
    return lines


def entries_bytes(entries):
    return "\n".join(entries).encode("utf-8")


def draw_params(draw_id, winners, entry_count, beacon_round=None):
    params = "id=%s\nwinners=%d\nentries=%d" % (draw_id, winners, entry_count)
    if beacon_round is not None:
        params += "\nbeacon_round=%d" % beacon_round
    return params


class SeedStream:
    """uint32 stream: HMAC-SHA256(final_seed, VERSION || uint32_be(i)), read big-endian."""

    def __init__(self, final_seed, version):
        self.key = final_seed
        self.label = version.encode("ascii")
        self.counter = 0
        self.buffer = b""

    def next_uint32(self):
        if not self.buffer:
            msg = self.label + struct.pack(">I", self.counter)
            self.buffer = hmac.new(self.key, msg, hashlib.sha256).digest()
            self.counter += 1
        value = struct.unpack(">I", self.buffer[:4])[0]
        self.buffer = self.buffer[4:]
        return value

    def uniform(self, m):
        limit = (2**32 // m) * m
        while True:
            u = self.next_uint32()
            if u < limit:
                return u % m


def compute_draw(draw_id, server_seed_hex, entries, winners, algorithm=V1, beacon=None):
    """beacon: {"round": int, "randomness": hex} for v2."""
    n = len(entries)
    if not 1 <= winners <= n:
        raise ValueError("winners must be between 1 and %d" % n)
    seed = bytes.fromhex(server_seed_hex)
    if len(seed) != 32:
        raise ValueError("server_seed must be 32 bytes")

    entries_hash = hashlib.sha256(entries_bytes(entries)).digest()
    prefix = algorithm.encode("ascii") + b"\x00" + seed + entries_hash
    if algorithm == V1:
        params = draw_params(draw_id, winners, n)
        final_seed = hashlib.sha256(prefix + params.encode("ascii")).digest()
    elif algorithm == V2:
        randomness = bytes.fromhex(beacon["randomness"])
        if len(randomness) != 32:
            raise ValueError("beacon randomness must be 32 bytes")
        params = draw_params(draw_id, winners, n, beacon["round"])
        final_seed = hashlib.sha256(prefix + randomness + params.encode("ascii")).digest()
    else:
        raise ValueError("unsupported algorithm %r" % algorithm)

    stream = SeedStream(final_seed, algorithm)
    a = list(range(n))
    for i in range(winners):
        j = i + stream.uniform(n - i)
        a[i], a[j] = a[j], a[i]

    return {
        "entries_hash": entries_hash.hex(),
        "commitment": hashlib.sha256(seed).hexdigest(),
        "params": params,
        "final_seed": final_seed.hex(),
        "winners": [
            {"rank": rank + 1, "index": index, "entry": entries[index]}
            for rank, index in enumerate(a[:winners])
        ],
    }


def lock_statement(r):
    beacon = r.get("beacon") or {}
    return "\n".join([
        "randnext-draw-lock-v1",
        "id=" + r["id"],
        "algorithm=" + r["algorithm"],
        "entries_hash=" + r["entries_hash"],
        "commitment=" + r["commitment"],
        "winners=%d" % r["winner_count"],
        "entries=%d" % r["entry_count"],
        "beacon_round=" + (str(beacon["round"]) if beacon.get("round") is not None else ""),
        "draw_at=" + r["draw_at"],
        "created_at=" + r["created_at"],
    ])


def result_statement(r):
    beacon = r.get("beacon") or {}
    return "\n".join([
        "randnext-draw-result-v1",
        "id=" + r["id"],
        "algorithm=" + r["algorithm"],
        "entries_hash=" + r["entries_hash"],
        "commitment=" + r["commitment"],
        "server_seed=" + r["server_seed"],
        "beacon_round=" + (str(beacon["round"]) if beacon.get("round") is not None else ""),
        "beacon_randomness=" + (beacon.get("randomness") or ""),
        "final_seed=" + r["final_seed"],
        "winners=" + ",".join(str(w["index"]) for w in r["winners"] or []),
        "drawn_at=" + r["drawn_at"],
    ])


def ed25519_verify(public_key_hex, signature_hex, message):
    """True/False, or None when the `cryptography` package isn't installed."""
    try:
        from cryptography.exceptions import InvalidSignature
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
    except ImportError:
        return None
    try:
        Ed25519PublicKey.from_public_bytes(bytes.fromhex(public_key_hex)).verify(bytes.fromhex(signature_hex), message.encode("utf-8"))
        return True
    except (InvalidSignature, ValueError):
        return False


def verify(record, entries_text, public_key=None, drand_beacon=None):
    """Returns ([(check name, ok), ...], recomputed result, [notes])."""
    if not record.get("server_seed"):
        raise SystemExit("This draw has not been revealed yet.")
    algorithm = record.get("algorithm")
    if algorithm not in (V1, V2):
        raise SystemExit("Unsupported algorithm %r (this script checks %s and %s)." % (algorithm, V1, V2))
    checks, notes = [], []

    beacon = None
    if algorithm == V2:
        b = record.get("beacon") or {}
        if b.get("round") is None or not b.get("randomness") or not b.get("signature"):
            return [("beacon", False)], None, ["The record has no revealed beacon."]
        checks.append(("beacon_randomness", hashlib.sha256(bytes.fromhex(b["signature"])).hexdigest() == b["randomness"]))
        if drand_beacon is not None:
            checks.append(("beacon_matches_drand", drand_beacon.get("randomness") == b["randomness"]))
        else:
            notes.append("Beacon not compared with drand (use --fetch, or check %s%d)." % (DRAND_QUICKNET, b["round"]))
        notes.append("drand's BLS signature is not checked here (no BLS in the standard library); the web verifier checks it.")
        beacon = {"round": b["round"], "randomness": b["randomness"]}

    entries = normalize_entries(entries_text, record.get("duplicates", "keep"))
    result = compute_draw(record["id"], record["server_seed"], entries, record["winner_count"], algorithm, beacon)
    checks += [
        ("commitment", result["commitment"] == record["commitment"]),
        ("entries_hash", result["entries_hash"] == record["entries_hash"]),
        ("entry_count", len(entries) == record["entry_count"]),
    ]
    if record.get("params") is not None:
        checks.append(("params", result["params"] == record["params"]))
    checks.append(("final_seed", result["final_seed"] == record["final_seed"]))
    # Records removed by moderation publish winner positions without names (entry is null).
    published = record["winners"] or []
    same_winners = len(published) == len(result["winners"]) and all(
        p["rank"] == w["rank"] and p["index"] == w["index"] and (p.get("entry") is None or p["entry"] == w["entry"])
        for p, w in zip(published, result["winners"])
    )
    checks.append(("winners", record["winners"] is not None and same_winners))

    sigs = record.get("signatures") or {}
    if sigs.get("lock") or sigs.get("result"):
        if not public_key:
            notes.append("Signatures not checked: no public key (use --fetch or --public-key).")
        else:
            for name, statement, sig in (("lock_signature", lock_statement, sigs.get("lock")), ("result_signature", result_statement, sigs.get("result"))):
                if not sig:
                    continue
                ok = ed25519_verify(public_key, sig, statement(record))
                if ok is None:
                    notes.append("Signatures not checked: pip install cryptography to check them.")
                    break
                checks.append((name, ok))
    return checks, result, notes


def fetch(url):
    with urllib.request.urlopen(url, timeout=30) as resp:
        return resp.read()


def run_vectors(path):
    with open(path, encoding="utf-8") as f:
        vectors = json.load(f)["vectors"]
    failures = 0
    for v in vectors:
        inp = v["input"]
        entries = normalize_entries(inp["entries_text"], inp["duplicates"])
        ok = entries == v["expected"]["entries"]
        if ok:
            result = compute_draw(inp["id"], inp["server_seed"], entries, inp["winners"], inp.get("algorithm", V1), inp.get("beacon"))
            expected = {k: v["expected"][k] for k in result}
            ok = result == expected
        failures += not ok
        print("%s  %s" % ("PASS" if ok else "FAIL", v["name"]))
    print("%d/%d vectors passed" % (len(vectors) - failures, len(vectors)))
    return failures == 0


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--fetch", metavar="DRAW_ID", help="download the record (and public entries) from randnext.com")
    group.add_argument("--record", metavar="FILE", help="record JSON from GET /api/v1/draws/{id}")
    group.add_argument("--vectors", metavar="FILE", help="check this implementation against test vectors")
    parser.add_argument("--entries", metavar="FILE", help="entry list, one per line")
    parser.add_argument("--public-key", metavar="HEX", help="RandNext's Ed25519 public key, for --record")
    args = parser.parse_args()

    if args.vectors:
        sys.exit(0 if run_vectors(args.vectors) else 1)

    public_key = args.public_key
    drand_beacon = None
    if args.fetch:
        record = json.loads(fetch(API_BASE + args.fetch))
        if not public_key and (record.get("signatures") or {}).get("key_id"):
            keys = json.loads(fetch(API_BASE + "signing-keys"))["keys"]
            match = [k for k in keys if k["key_id"] == record["signatures"]["key_id"]]
            public_key = match[0]["public_key"] if match else None
        beacon = record.get("beacon") or {}
        if beacon.get("round") is not None and beacon.get("randomness"):
            drand_beacon = json.loads(fetch(DRAND_QUICKNET + str(beacon["round"])))
    else:
        if not args.entries:
            parser.error("--record needs --entries")
        with open(args.record, encoding="utf-8") as f:
            record = json.load(f)

    if args.entries:
        # newline="" keeps the bytes as published; normalization handles \r itself.
        with open(args.entries, encoding="utf-8", newline="") as f:
            entries_text = f.read()
    elif record.get("entries_url"):
        entries_text = fetch(record["entries_url"]).decode("utf-8")
    else:
        sys.exit("This draw's entries are not public. Pass the list with --entries.")

    checks, result, notes = verify(record, entries_text, public_key, drand_beacon)
    for name, ok in checks:
        print("%s  %s" % ("PASS" if ok else "FAIL", name))
    for note in notes:
        print("NOTE  " + note)
    if result:
        print("\nRecomputed winners:")
        for w in result["winners"]:
            print("#%d  %s  (entry %d)" % (w["rank"], w["entry"], w["index"]))
    sys.exit(0 if checks and all(ok for _, ok in checks) else 1)


if __name__ == "__main__":
    main()