#!/usr/bin/env python3
"""Offline verifier for a PolicyCortex Evidence_History/ledger.jsonl stream.

Usage:
  python3 verify_ledger.py [ledger.jsonl] [--organization-id ORG] [--ledger ledger.json] [--key-file KEY]

Standard library only (Python 3.6 or newer). Exit code 0 when the chain is
INTACT and agrees with ledger.json, 1 when it is BROKEN or disagrees, 2 when the
file cannot be read.

What is checked, for every line in order:
  1. The line is a JSON object whose record_type is state_capture,
     change_record, declared_gap or agency_record, with the fields that type
     requires (hex digests where a digest is required, null where a state or
     rollback id was not observed).
  2. sequence equals the line position, line 0 has previousDigest null, and
     every later line has previousDigest equal to the digest of the line before.
  3. digest equals SHA3-256 over the canonical JSON of the record without its
     digest, id and sequence fields (previousDigest lowercased or null).
  4. With --organization-id, every record belongs to that organization.
When ledger.json sits beside the stream (or --ledger names it), its
sha3Chain.headDigest and record counts must match the stream, and a signed
ledger.json must carry a signedPayload that binds this stream: after the
SHA-256 rootHash it lists sha3Chain=2, headDigest and the four record counts,
which must equal what the stream recomputes to. With --key-file (a file holding
the operator's signing key, as issued), the HMAC-SHA256 signature over that
payload is checked too.
"""
import hashlib
import hmac
import json
import os
import re
import sys

RECORD_TYPES = ("state_capture", "change_record", "declared_gap", "agency_record")
HEX64 = re.compile(r"^[0-9a-fA-F]{64}$")
INDEX_KEY = re.compile(r"^(0|[1-9][0-9]*)$")
TIMESTAMP = re.compile(
    r"^(\d{4})-(\d{2})-(\d{2})(?:T(\d{2}):(\d{2})(?::(\d{2})(?:\.(\d+))?)?(Z|[+-]\d{2}:\d{2})?)?$"
)


class RawNumber(str):
    """A JSON number token kept exactly as JSON.stringify wrote it."""


class SchemaError(Exception):
    pass


def reject_constant(name):
    raise ValueError("non-JSON constant " + name)


def parse_line(line):
    return json.loads(line, parse_float=RawNumber, parse_int=RawNumber, parse_constant=reject_constant)


def js_string(value):
    text = json.dumps(value, ensure_ascii=False)
    # JSON.stringify escapes lone surrogates; json.dumps leaves them raw.
    return "".join("\\u%04x" % ord(ch) if 0xD800 <= ord(ch) <= 0xDFFF else ch for ch in text)


def js_key_order(keys):
    keys = list(keys)
    index = sorted((k for k in keys if INDEX_KEY.match(k) and int(k) <= 4294967294), key=int)
    taken = set(index)
    rest = sorted((k for k in keys if k not in taken), key=lambda k: k.encode("utf-16-be", "surrogatepass"))
    return index + rest


def canonical(value):
    if isinstance(value, RawNumber):
        return str(value)
    if value is None:
        return "null"
    if value is True:
        return "true"
    if value is False:
        return "false"
    if isinstance(value, str):
        return js_string(value)
    if isinstance(value, list):
        return "[" + ",".join(canonical(item) for item in value) + "]"
    if isinstance(value, dict):
        return "{" + ",".join(js_string(k) + ":" + canonical(value[k]) for k in js_key_order(value.keys())) + "}"
    raise SchemaError("value is not JSON")


def record_digest(record, previous_digest):
    material = {k: v for k, v in record.items() if k not in ("digest", "id", "sequence")}
    material["previousDigest"] = previous_digest.lower() if isinstance(previous_digest, str) and previous_digest else None
    return hashlib.sha3_256(canonical(material).encode("utf-8", "surrogatepass")).hexdigest()


def is_hex(value):
    return isinstance(value, str) and HEX64.match(value) is not None


def timestamp_key(value, field):
    match = TIMESTAMP.match(value) if isinstance(value, str) else None
    if not match:
        raise SchemaError(field + " is not an ISO-8601 timestamp")
    year, month, day, hour, minute, second, fraction, zone = match.groups()
    offset = 0
    if zone and zone != "Z":
        sign = 1 if zone[0] == "+" else -1
        offset = sign * (int(zone[1:3]) * 60 + int(zone[4:6]))
    minutes = (int(hour or 0) * 60 + int(minute or 0)) - offset
    return (int(year), int(month), int(day), minutes, int(second or 0), int((fraction or "0")[:3].ljust(3, "0")))


def require_text(record, field, label):
    value = record.get(field)
    if not isinstance(value, str) or not value.strip():
        raise SchemaError(label + " " + field + " is required")


def require_hex(record, field, label, nullable=False):
    value = record.get(field)
    if nullable and value is None:
        return
    if not is_hex(value):
        raise SchemaError(label + " " + field + " is invalid")


def require_rollback_id(record, label):
    value = record.get("rollbackId")
    if value is not None and (not isinstance(value, str) or not value.strip()):
        raise SchemaError(label + " rollbackId must be a non-empty string or null")


def validate_schema(record):
    kind = record.get("record_type")
    if kind not in RECORD_TYPES:
        raise SchemaError("Unsupported stream record_type: " + str(kind))
    previous = record.get("previousDigest")
    if previous is not None and not is_hex(previous):
        raise SchemaError(kind + " previousDigest is invalid")
    if kind == "state_capture":
        require_hex(record, "contentHash", "State record")
        size = record.get("bodySize")
        if not isinstance(size, RawNumber) or not re.match(r"^(0|[1-9][0-9]*)$", size):
            raise SchemaError("State record bodySize is invalid")
        timestamp_key(record.get("capturedAt"), "capturedAt")
    elif kind == "change_record":
        for field in ("beforeStateHash", "afterStateHash", "deltaHash"):
            require_hex(record, field, "Change record")
        require_rollback_id(record, "Change record")
        timestamp_key(record.get("capturedAt"), "capturedAt")
    elif kind == "declared_gap":
        for field in ("organizationId", "stream", "reason"):
            require_text(record, field, "Declared gap")
        if timestamp_key(record.get("from"), "from") > timestamp_key(record.get("to"), "to"):
            raise SchemaError("Declared gap interval from must not exceed to")
        timestamp_key(record.get("declared_at"), "declared_at")
    else:
        for field in ("organizationId", "stream"):
            require_text(record, field, "Agency record")
        require_hex(record, "envelopeVersionDigest", "Agency record")
        decision = record.get("policyDecision")
        if decision is None or (isinstance(decision, str) and not decision.strip()):
            raise SchemaError("Agency record policyDecision is required")
        require_hex(record, "beforeStateHash", "Agency record", nullable=True)
        require_hex(record, "afterStateHash", "Agency record", nullable=True)
        if "deltaHash" in record:
            require_hex(record, "deltaHash", "Agency record")
        require_rollback_id(record, "Agency record")
        require_hex(record, "counterfactualDigest", "Agency record")
        timestamp_key(record.get("capturedAt"), "capturedAt")


def verify_stream(lines, organization_id=None):
    counts = dict((kind, 0) for kind in RECORD_TYPES)
    previous = None
    for index, line in enumerate(lines):
        record_id = "line-" + str(index)
        try:
            record = parse_line(line)
        except ValueError:
            return {"status": "BROKEN", "firstBrokenSequence": index, "brokenRecordId": record_id,
                    "reason": "Malformed JSON at line " + str(index), "counts": counts}
        if not isinstance(record, dict):
            return {"status": "BROKEN", "firstBrokenSequence": index, "brokenRecordId": record_id,
                    "reason": "Stream line " + str(index) + " is not a JSON object", "counts": counts}
        if isinstance(record.get("id"), str) and record["id"].strip():
            record_id = record["id"]

        def broken(reason):
            return {"status": "BROKEN", "firstBrokenSequence": index, "brokenRecordId": record_id,
                    "reason": reason, "counts": counts}

        sequence = record.get("sequence")
        if sequence is not None and str(sequence) != str(index):
            return broken("Sequence mismatch at line " + str(index) + ": record says " + str(sequence))
        if organization_id is not None and record.get("organizationId") != organization_id:
            return broken("Tenant isolation violation at sequence " + str(index) + ": expected organizationId " + organization_id)
        try:
            validate_schema(record)
        except SchemaError as error:
            return broken(str(error))
        stored = record.get("digest")
        if not is_hex(stored):
            return broken("Missing or invalid digest at sequence " + str(index))
        claimed = record.get("previousDigest")
        claimed = claimed.lower() if isinstance(claimed, str) else None
        if index == 0 and claimed is not None:
            return broken("Genesis record at sequence 0 must have previousDigest null")
        if index > 0 and claimed != previous:
            return broken("Previous digest mismatch at sequence " + str(index) + ": expected " + str(previous) + ", got " + str(claimed))
        recomputed = record_digest(record, claimed)
        if recomputed != stored.lower():
            return broken("Digest mismatch at sequence " + str(index) + ": stored " + stored.lower() + " does not match recomputed " + recomputed)
        previous = recomputed
        counts[record["record_type"]] += 1
    return {"status": "INTACT", "headDigest": previous, "recordCount": len(lines), "counts": counts}


LEDGER_COUNT_FIELDS = (
    ("totalRecords", None),
    ("stateCaptureCount", "state_capture"),
    ("changeRecordCount", "change_record"),
    ("agencyRecordCount", "agency_record"),
    ("declaredGapCount", "declared_gap"),
)


SIGNED_STREAM_FIELDS = (
    ("stateCaptureCount", "state_capture"),
    ("changeRecordCount", "change_record"),
    ("agencyRecordCount", "agency_record"),
    ("declaredGapCount", "declared_gap"),
)


def read_key(key_path):
    with open(key_path, "rb") as handle:
        key = handle.read()
    for ending in (b"\r\n", b"\n"):
        if key.endswith(ending):
            return key[: -len(ending)]
    return key


def check_signature(ledger, result, key_path):
    """Returns (problems, notes) for the ledger.json signature."""
    signature = ledger.get("signature")
    if not signature:
        return [], ["signature: ledger.json is not signed (" + str(ledger.get("signatureStatus")) + "); keep the head digest outside PolicyCortex to detect later changes"]
    payload = signature.get("signedPayload")
    fields = payload.split("|") if isinstance(payload, str) else []
    if len(fields) < 7 or fields[0] != "pcx-evidence-ledger":
        return ["ledger.json signature has no PolicyCortex evidence ledger signedPayload"], []
    problems = []
    root = (ledger.get("verification") or {}).get("rootHash")
    if isinstance(root, str) and fields[6] != root.lower():
        problems.append("signedPayload rootHash " + fields[6] + " does not match ledger.json verification.rootHash " + root)
    binding = dict(field.split("=", 1) for field in fields[7:] if "=" in field)
    if binding.get("sha3Chain") != "2":
        problems.append("the signature does not bind the SHA3-256 stream: signedPayload has no sha3Chain=2 binding")
    else:
        expected = [("headDigest", result["headDigest"] or "null")]
        expected += [(field, str(result["counts"][kind])) for field, kind in SIGNED_STREAM_FIELDS]
        for field, value in expected:
            if binding.get(field) != value:
                problems.append("signedPayload " + field + " is " + str(binding.get(field)) + " but the stream has " + value)
    notes = []
    if key_path:
        key = read_key(key_path)
        if hashlib.sha256(key).hexdigest()[:16] != signature.get("keyId"):
            problems.append("the supplied key is not key id " + str(signature.get("keyId")))
        else:
            mac = hmac.new(key, payload.encode("utf-8"), hashlib.sha256).hexdigest()
            if not hmac.compare_digest(mac, str(signature.get("signature")).lower()):
                problems.append("HMAC-SHA256 signature does not verify over signedPayload with the supplied key")
            else:
                notes.append("signature: HMAC-SHA256 verified with key id " + str(signature.get("keyId")) + "; it binds the root, the stream head and the record counts")
    elif not problems:
        notes.append("signature: signedPayload binds the stream head and record counts; pass --key-file to check the HMAC with key id " + str(signature.get("keyId")))
    return problems, notes


def compare_with_ledger(result, ledger_path, key_path=None):
    with open(ledger_path, "r", encoding="utf-8") as handle:
        ledger = json.load(handle) or {}
    chain = ledger.get("sha3Chain") or {}
    problems = []
    if chain.get("headDigest") != result["headDigest"]:
        problems.append("ledger.json sha3Chain.headDigest " + str(chain.get("headDigest")) + " does not match stream head " + str(result["headDigest"]))
    for field, kind in LEDGER_COUNT_FIELDS:
        if field not in chain:
            continue
        actual = result["recordCount"] if kind is None else result["counts"][kind]
        if chain[field] != actual:
            problems.append("ledger.json sha3Chain." + field + " is " + str(chain[field]) + " but the stream has " + str(actual))
    signature_problems, notes = check_signature(ledger, result, key_path)
    return chain, problems + signature_problems, notes


def main(argv):
    args = list(argv)
    organization_id = None
    ledger_path = None
    key_path = None
    if "--key-file" in args:
        position = args.index("--key-file")
        key_path = args[position + 1] if position + 1 < len(args) else None
        del args[position:position + 2]
    if "--organization-id" in args:
        position = args.index("--organization-id")
        organization_id = args[position + 1] if position + 1 < len(args) else None
        del args[position:position + 2]
    if "--ledger" in args:
        position = args.index("--ledger")
        ledger_path = args[position + 1] if position + 1 < len(args) else None
        del args[position:position + 2]
    stream_path = args[0] if args else os.path.join(os.path.dirname(os.path.abspath(__file__)), "ledger.jsonl")
    if ledger_path is None:
        candidate = os.path.join(os.path.dirname(os.path.abspath(stream_path)), "ledger.json")
        ledger_path = candidate if os.path.exists(candidate) else None
    try:
        with open(stream_path, "r", encoding="utf-8", newline="") as handle:
            # Split on newline only: str.splitlines() also splits on U+2028 and
            # U+2029, which JSON.stringify leaves unescaped inside strings.
            lines = [line.strip(" \t\r") for line in handle.read().split("\n") if line.strip(" \t\r")]
    except OSError as error:
        print("Cannot read " + stream_path + ": " + str(error))
        return 2

    result = verify_stream(lines, organization_id)
    counts = result["counts"]
    print("PolicyCortex evidence stream verification (SHA3-256)")
    print("stream: " + stream_path)
    print("status: " + result["status"])
    if result["status"] == "BROKEN":
        print("firstBrokenSequence: " + str(result["firstBrokenSequence"]))
        print("brokenRecordId: " + str(result["brokenRecordId"]))
        print("reason: " + result["reason"])
        return 1
    print("headDigest: " + str(result["headDigest"]))
    print("records: " + str(result["recordCount"]) + " (" + ", ".join(kind + " " + str(counts[kind]) for kind in RECORD_TYPES) + ")")
    if key_path and not ledger_path:
        print("MISMATCH: --key-file needs ledger.json beside the stream or --ledger")
        return 1
    if ledger_path:
        try:
            chain, problems, notes = compare_with_ledger(result, ledger_path, key_path)
        except (OSError, ValueError) as error:
            print("Cannot read " + ledger_path + ": " + str(error))
            return 2
        if problems:
            for problem in problems:
                print("MISMATCH: " + problem)
            return 1
        print("ledger.json: sha3Chain head digest and counts match")
        for note in notes:
            print(note)
        if "changesWithoutObservedState" in chain or "agencyDecisionsWithoutEnvelope" in chain:
            print("not chained: " + str(chain.get("changesWithoutObservedState", 0)) + " executed changes without an observed before and after state, "
                  + str(chain.get("agencyDecisionsWithoutEnvelope", 0)) + " autonomous decisions without a recorded envelope evaluation (ids in manifest.json evidenceChainV2.exclusions)")
    return 0


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