#!/usr/bin/env python3
"""Verify a Golem evidence export offline, from the zip alone.

    pip install cbor2 cryptography
    python3 verify.py golem-evidence-....zip

Checks, and prints a line for each:
  1. every file the manifest lists is present and hashes to what it says;
  2. the manifest's COSE_Sign1 signature, against the evidence key it names;
  3. the workspace evidence chain: each entry links to the one before it,
     recomputes to its entry_hash, and is signed by the key it names;
  4. the entry recording this export follows the chain and commits to the
     manifest's hash;
  5. each ledger session chain the export carries: links, hashes, and
     signatures against the hosts' public keys;
  6. the population reconciles: raised = resolved + escalated + open;
  7. the Merkle anchor: the root it names commits to the evidence chain tip
     it names, and that tip is an entry of the exported chain. Compare the
     root with the copy the deployment published write-once.

Exits 0 when everything verifies, 1 otherwise. Tamper-evident, not
tamper-proof: a pass means nothing in the package was changed after Golem
signed it, by anyone without the signing keys.
"""

import base64
import csv
import hashlib
import io
import json
import sys
import zipfile

import cbor2
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature

GENESIS = "0" * 64


class Report:
    def __init__(self):
        self.failed = 0

    def ok(self, msg):
        print(f"  ok    {msg}")

    def fail(self, msg):
        self.failed += 1
        print(f"  FAIL  {msg}")

    def note(self, msg):
        print(f"  note  {msg}")


def sha256_hex(b):
    return hashlib.sha256(b).hexdigest()


def entry_hash(payload, previous):
    return sha256_hex(payload + previous.encode("utf-8"))


def cose_verify(cose_bytes, public_key):
    """Return (payload, kid) when the COSE_Sign1 signature verifies, else raise."""
    msg = cbor2.loads(cose_bytes)
    if isinstance(msg, cbor2.CBORTag):
        msg = msg.value
    protected, _unprotected, payload, signature = msg
    headers = cbor2.loads(protected) if protected else {}
    if headers.get(1) != -7:
        raise ValueError("not ES256")
    to_be_signed = cbor2.dumps(["Signature1", protected, b"", payload])
    if len(signature) != 64:
        raise ValueError("signature is not a 64-byte ES256 signature")
    der = encode_dss_signature(int.from_bytes(signature[:32], "big"), int.from_bytes(signature[32:], "big"))
    public_key.verify(der, to_be_signed, ec.ECDSA(hashes.SHA256()))
    kid = headers.get(4, b"")
    return payload, kid.decode() if isinstance(kid, bytes) else str(kid)


def cose_kid(cose_bytes):
    msg = cbor2.loads(cose_bytes)
    if isinstance(msg, cbor2.CBORTag):
        msg = msg.value
    headers = cbor2.loads(msg[0]) if msg[0] else {}
    kid = headers.get(4, b"")
    return kid.decode() if isinstance(kid, bytes) else str(kid)


def read_jsonl(z, name):
    return [json.loads(line) for line in z.read(name).decode().splitlines() if line.strip()]


def verify_chain(r, label, entries, keys, start_previous):
    prev = start_previous
    bad = 0
    unsigned = 0
    for e in entries:
        payload = base64.b64decode(e["payload_cbor"])
        if e["previous_hash"] != prev:
            r.fail(f"{label}: entry {e.get('seq')} does not follow the one before it")
            bad += 1
        if entry_hash(payload, e["previous_hash"]) != e["entry_hash"]:
            r.fail(f"{label}: entry {e.get('seq')} does not recompute to its entry_hash")
            bad += 1
        cose = base64.b64decode(e.get("cose_sign1") or "")
        if cose:
            kid = cose_kid(cose)
            key = keys.get(kid)
            if key is None:
                unsigned += 1
            else:
                try:
                    signed, _ = cose_verify(cose, key)
                    if signed != payload:
                        r.fail(f"{label}: entry {e.get('seq')} is signed over something other than its payload")
                        bad += 1
                except (InvalidSignature, ValueError):
                    r.fail(f"{label}: entry {e.get('seq')} has a bad signature")
                    bad += 1
        else:
            unsigned += 1
        prev = e["entry_hash"]
    if bad == 0:
        r.ok(f"{label}: {len(entries)} entries link, recompute and verify")
    if unsigned:
        r.note(f"{label}: {unsigned} entries carry no signature, or a key this export does not include")
    return prev


def merkle_leaf(tip):
    data = tip["s"].encode() + b"\x1f" + tip["h"].encode() + b"\x1f" + str(tip["q"]).encode()
    return hashlib.sha256(b"\x00" + data).digest()


def merkle_node(left, right):
    return hashlib.sha256(b"\x01" + left + right).digest()


def verify_inclusion(root, leaf, index, size, proof):
    if index < 0 or size <= 0 or index >= size:
        return False
    inner = (index ^ (size - 1)).bit_length()
    border = bin(index >> inner).count("1")
    if len(proof) != inner + border:
        return False
    res = leaf
    for i, h in enumerate(proof[:inner]):
        res = merkle_node(res, h) if (index >> i) & 1 == 0 else merkle_node(h, res)
    for h in proof[inner:]:
        res = merkle_node(h, res)
    return res == root


def load_pem(pem):
    return serialization.load_pem_public_key(pem.encode())


def load_der(b64):
    return serialization.load_der_public_key(base64.b64decode(b64))


def main(path):
    r = Report()
    z = zipfile.ZipFile(path)
    names = set(z.namelist())
    print(f"Verifying {path}")

    manifest_bytes = z.read("manifest.json")
    manifest = json.loads(manifest_bytes)
    for f in manifest["files"]:
        if f["path"] not in names:
            r.fail(f"{f['path']} is listed but missing")
        elif sha256_hex(z.read(f["path"])) != f["sha256"]:
            r.fail(f"{f['path']} does not hash to what the manifest says")
    if r.failed == 0:
        r.ok(f"{len(manifest['files'])} files match the manifest")

    evidence_keys = {}
    for k in json.loads(z.read("keys/evidence.json")):
        try:
            evidence_keys[k["kid"]] = load_pem(k["public_key_pem"])
        except ValueError:
            r.note(f"evidence key {k['kid'][:16]}… could not be parsed; entries it signed will not verify")
    try:
        signed, kid = cose_verify(z.read("manifest.cose"), evidence_keys[cose_kid(z.read("manifest.cose"))])
        if signed == manifest_bytes:
            r.ok(f"manifest signature verifies (key {kid[:16]}…)")
        else:
            r.fail("the manifest signature covers different bytes")
    except (KeyError, InvalidSignature, ValueError):
        r.fail("the manifest signature does not verify")

    chain = read_jsonl(z, "evidence/chain.jsonl")
    start = manifest.get("evidence_chain", {}).get("starts_after", GENESIS)
    if start != GENESIS:
        r.note(f"the evidence chain is a segment starting after {start[:16]}…")
    tip = verify_chain(r, "evidence chain", chain, evidence_keys, start)

    export_entry = json.loads(z.read("export_entry.json"))
    tip = verify_chain(r, "export entry", [export_entry], evidence_keys, tip)
    detail = cbor2.loads(base64.b64decode(export_entry["payload_cbor"])).get("detail", {})
    if detail.get("manifest_sha256") == sha256_hex(manifest_bytes):
        r.ok("the export's own evidence entry commits to this manifest")
    else:
        r.fail("the export's evidence entry names a different manifest")

    host_keys = {}
    if "keys/hosts.json" in names:
        for k in json.loads(z.read("keys/hosts.json")):
            try:
                host_keys[k["key_id"]] = load_der(k["public_key_der"])
            except ValueError:
                r.note(f"host key {k['key_id'][:16]}… could not be parsed")
    bom_payload = {}
    for name in sorted(n for n in names if n.startswith("ledger/") and n.endswith(".jsonl")):
        entries = read_jsonl(z, name)
        if entries:
            verify_chain(r, name, entries, host_keys, entries[0]["previous_hash"])
        for e in entries:
            try:
                bom = cbor2.loads(base64.b64decode(e["payload_cbor"]))
            except Exception as err:
                r.note(f"{name}: an entry's BOM could not be decoded ({type(err).__name__})")
                continue
            if isinstance(bom, dict) and bom.get("payload_hash"):
                bom_payload[e["entry_hash"]] = bom["payload_hash"]

    if "payloads/index.json" in names:
        index = json.loads(z.read("payloads/index.json"))["entries"]
        present, missing = 0, 0
        for item in index:
            if bom_payload.get(item["entry_hash"]) != item["payload_hash"]:
                r.fail(f"payloads/index.json names {item['payload_hash'][:16]}… for an entry whose signed BOM says otherwise")
                continue
            if item["file"] not in names:
                missing += 1
                continue
            if sha256_hex(z.read(item["file"])) == item["payload_hash"]:
                present += 1
            else:
                r.fail(f"{item['file']} does not hash to the payload_hash its signed entry commits to")
        if present:
            r.ok(f"{present} payload bodies match the hashes their signed ledger entries commit to")
        if missing:
            r.note(f"{missing} indexed payloads were not included in this package")

    if "reconciliation.json" in names:
        rec = json.loads(z.read("reconciliation.json"))
        parts = rec["resolved"] + rec["escalated_open"] + rec["open_awaiting_triage"]
        if rec["raised"] == parts:
            r.ok(f"population reconciles: {rec['raised']} raised = {rec['resolved']} resolved + "
                 f"{rec['escalated_open']} in open cases + {rec['open_awaiting_triage']} awaiting triage")
        else:
            r.fail(f"population does not reconcile: {rec['raised']} raised, {parts} accounted for")
        rows = list(csv.DictReader(io.StringIO(z.read("population.csv").decode())))
        if len(rows) == rec["raised"]:
            r.ok(f"population.csv lists all {len(rows)} flags")
        else:
            r.fail(f"population.csv lists {len(rows)} flags, the reconciliation counts {rec['raised']}")

    if "anchor.json" in names:
        anchor = json.loads(z.read("anchor.json"))
        if anchor.get("status") != "anchored":
            r.note(f"the evidence chain is not anchored in a Merkle root yet ({anchor.get('status')})")
        else:
            leaf = anchor["leaf"]
            proof = [bytes.fromhex(p) for p in anchor["proof"]]
            if verify_inclusion(bytes.fromhex(anchor["root_hash"]), merkle_leaf(leaf), anchor["leaf_index"],
                                anchor["leaf_count"], proof):
                r.ok(f"Merkle root {anchor['root_hash'][:16]}… of {anchor['computed_at']} commits to the evidence chain at seq {leaf['q']}")
            else:
                r.fail("the Merkle inclusion proof does not verify")
            if any(e["entry_hash"] == leaf["h"] and e["seq"] == leaf["q"] for e in chain):
                r.ok("the anchored tip is an entry of the exported chain")
            elif chain and chain[0]["seq"] > leaf["q"]:
                r.note("the anchored tip is older than the exported chain segment")
            else:
                r.fail("the anchored tip is not in the exported chain")
            r.note(f"compare root_hash with the published copy {anchor.get('published_as')}")

    print("VERIFIED" if r.failed == 0 else f"FAILED ({r.failed} problems)")
    return 0 if r.failed == 0 else 1


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(__doc__)
        sys.exit(2)
    sys.exit(main(sys.argv[1]))