"""NC-VER-0001. Check one document against the published root, or rebuild it.

This file holds no root, no hash and no path. Everything it compares against is
supplied on the command line by the reader, from the values NC-VER-0001
publishes. It makes no network request: fetch the bytes yourself, with whatever
you trust to fetch them, and hand them to this.

Standard library only. Python 3.6 or later, for SHA3-256 in hashlib.

  One document, against the published proof. Every value comes off the page;
  none of them is written here. An earlier draft of this paragraph did write
  them, with four sibling placeholders against an index and a count that need
  five, so the one worked shape in the file that asks to be read rather than
  believed was a shape that fails:

    python verify.py note-0147.html \\
        --index <n> --count <leaves> \\
        --root <root hex> \\
        --path <hex>,<hex>,...

  The whole tree, if you would rather not trust the path either. Fetch every
  URL in sitemap.xml into one directory, named for its route with .html
  appended, the front page as index.html:

    python verify.py --all ./leaves --sitemap sitemap.xml --root <root hex>

Construction, so that this file can be read instead of believed:

    leaf   SHA3-256( 0x00 || document bytes )
    node   SHA3-256( 0x01 || left || right )

The two prefixes are the domain separation of RFC 6962 section 2.1: without
them a leaf could be presented as an interior node. An odd node at any level is
promoted to the next level unhashed rather than paired with itself, so the leaf
count is needed to walk the path -- the path length alone does not fix the
shape of the tree.

A pass means those bytes are in that root. It means nothing else. NC-VER-0001
sets out what it does not establish, and that list is longer than this one.
"""
import argparse
import hashlib
import re
import sys
from pathlib import Path


def leaf_hash(data):
    return hashlib.sha3_256(b"\x00" + data).digest()


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


def climb(leaf, index, count, path):
    """Walk one leaf to a root, consuming the path as the shape demands."""
    here, idx, width, steps = leaf, index, count, list(path)
    while width > 1:
        promoted = idx == width - 1 and width % 2 == 1
        if not promoted:
            if not steps:
                raise ValueError(
                    f"the path ran out at width {width}; a tree of {count} "
                    "leaves needs more siblings than were given")
            sibling = steps.pop(0)
            here = (node_hash(sibling, here) if idx % 2
                    else node_hash(here, sibling))
        idx //= 2
        width = (width + 1) // 2
    if steps:
        raise ValueError(f"{len(steps)} sibling(s) left over; the path is "
                         f"longer than a tree of {count} leaves has room for")
    return here


def whole_tree(directory, sitemap, not_a_leaf):
    routes = [re.sub(r"^https?://[^/]+/?", "", url)
              for url in re.findall(r"<loc>([^<]+)</loc>",
                                    Path(sitemap).read_text(encoding="utf-8"))]
    routes = [r for r in routes if r not in not_a_leaf]
    row = []
    for route in routes:
        name = "index.html" if route == "" else f"{route}.html"
        candidate = Path(directory) / name
        if not candidate.is_file():
            raise SystemExit(f"missing leaf: {candidate}")
        row.append(leaf_hash(candidate.read_bytes()))
    print(f"{len(row)} leaves read from {directory}")
    while len(row) > 1:
        nxt = [node_hash(row[i], row[i + 1]) for i in range(0, len(row) - 1, 2)]
        if len(row) % 2:
            nxt.append(row[-1])
        row = nxt
    return row[0]


def main(argv):
    ap = argparse.ArgumentParser(
        description="Check a document against the NC-VER-0001 root.")
    ap.add_argument("document", nargs="?",
                    help="the file whose bytes you fetched")
    ap.add_argument("--root", required=True, help="root hex, from NC-VER-0001")
    ap.add_argument("--index", type=int, help="the leaf's position, from 0")
    ap.add_argument("--count", type=int, help="how many leaves the tree has")
    ap.add_argument("--path", default="",
                    help="sibling hashes, comma separated, leaf upward")
    ap.add_argument("--all", metavar="DIR",
                    help="rebuild the root from a directory of every leaf")
    ap.add_argument("--sitemap", help="sitemap.xml, for --all")
    ap.add_argument("--not-a-leaf", default="verify",
                    help="routes in the sitemap that are not leaves, comma "
                         "separated (default: verify)")
    args = ap.parse_args(argv)

    expected = args.root.strip().lower()
    if not re.fullmatch(r"[0-9a-f]{64}", expected):
        print("the root must be 64 hex characters")
        return 2
    excluded = {r for r in args.not_a_leaf.split(",") if r}

    if args.all:
        if not args.sitemap:
            print("--all needs --sitemap")
            return 2
        got = whole_tree(args.all, args.sitemap, excluded).hex()
        print(f"computed root  {got}")
        print(f"published root {expected}")
        print("PASS: the leaves you fetched build the published root."
              if got == expected else
              "FAIL: they do not. One of the documents is not the one that was "
              "hashed, or the set is not the same set.")
        return 0 if got == expected else 1

    missing = [name for name, value in
               (("document", args.document), ("--index", args.index),
                ("--count", args.count)) if value is None]
    if missing:
        print(f"missing {', '.join(missing)}; see --help")
        return 2
    if not 0 <= args.index < args.count:
        print(f"--index {args.index} is not a leaf of a tree with "
              f"{args.count} of them")
        return 2

    data = Path(args.document).read_bytes()
    leaf = leaf_hash(data)
    steps = []
    for chunk in args.path.replace(" ", "").split(","):
        if not chunk:
            continue
        if not re.fullmatch(r"[0-9a-f]{64}", chunk.lower()):
            print(f"not a hash: {chunk}")
            return 2
        steps.append(bytes.fromhex(chunk))

    print(f"document    {args.document}, {len(data)} bytes")
    print(f"leaf hash   {leaf.hex()}")
    try:
        got = climb(leaf, args.index, args.count, steps).hex()
    except ValueError as err:
        print(f"FAIL: {err}")
        return 1
    print(f"computed    {got}")
    print(f"published   {expected}")
    if got == expected:
        print(f"\nPASS: these bytes are leaf {args.index} of that root.")
        print("That is all it establishes. Read section 5 before you rely on "
              "it for anything else.")
        return 0
    print("\nFAIL: these bytes are not leaf "
          f"{args.index} of that root. Either the document changed after the "
          "root was cut, or the path does not belong to this leaf.")
    return 1


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