#!/usr/bin/env python3
"""Verify exact sample artifact bytes against an independently trusted key.

Requires Python and cryptography. Trust arguments must come from a reviewed
source outside the downloaded bundle. A valid signature is not scan evidence.

When the bundle carries transparency.json (ADR 0090), the entry is checked
offline too: it must name this statement's digest, carry an ECDSA P-256
signature over the same statement bytes by the signing key's published
log-signing companion (supplied as a second trusted value, since the public
log can verify a signature only over a digest and a plain Ed25519 signature
does not allow that), its leaf hash must be the entry identifier, the
inclusion proof must reach the checkpoint root, and the checkpoint must be
signed by the pinned log key.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import json
from pathlib import Path
import re
import sys
import uuid

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey

# The transparency log the platform writes to (ADR 0090): log id = SHA-256
# of the key's DER encoding; the same copy is pinned in checkii.ai/verify.
KNOWN_LOGS = {
    "c0d23d6ad406973f9559f3ba2d1ca01f84147d8ffc5b8445c224f98b9591801d": {
        "origin": "rekor.sigstore.dev",
        "log_url": "https://rekor.sigstore.dev",
        "public_key": (
            "MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE2G2Y+2tabdTV5BcGiBIx0a9fAFwrkBbmLSGtks4L3qX6"
            "yYY0zufBnhC8Ur/iy55GhWP/9A/bY2LhC30M9+RYtw=="
        ),
    },
}

STATEMENT_FIELDS = {
    "schema_version",
    "report_id",
    "tenant_id",
    "engagement_id",
    "version",
    "report_type",
    "status",
    "title",
    "snapshot_sha256",
    "markdown_sha256",
    "created_by",
    "approved_by",
    "issued_at",
    "supersedes_id",
}


def canonical_bytes(payload: dict) -> bytes:
    return json.dumps(
        payload,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
        allow_nan=False,
    ).encode()


def bounded_read(path: Path, limit: int) -> bytes:
    with path.open("rb") as stream:
        raw = stream.read(limit + 1)
    if len(raw) > limit:
        raise ValueError("artifact size limit")
    return raw


def parse(raw: bytes):
    def pairs(items):
        result = {}
        for key, value in items:
            if key in result:
                raise ValueError("duplicate JSON key")
            result[key] = value
        return result

    def invalid_constant(value):
        raise ValueError("nonfinite JSON value")

    return json.loads(raw, object_pairs_hook=pairs, parse_constant=invalid_constant)


def verify(directory: Path, trusted_key_id: str, trusted_public_key: str) -> dict:
    bundle = parse(bounded_read(directory / "signature.json", 65536))
    if not re.fullmatch(r"[A-Za-z0-9._-]{1,100}", trusted_key_id):
        raise ValueError("invalid trusted key identifier")
    if (
        bundle["signature_algorithm"] != "Ed25519"
        or bundle["signing_key_id"] != trusted_key_id
        or bundle["signing_public_key"] != trusted_public_key
    ):
        raise ValueError("bundle key is not the independently trusted key")
    public_key = base64.b64decode(trusted_public_key, validate=True)
    statement = bundle["signed_statement"]
    if not isinstance(statement, dict) or set(statement) != STATEMENT_FIELDS:
        raise ValueError("incomplete or unexpected issuance statement")
    if (
        statement["schema_version"] != "aegis.report-signature.v1"
        or statement["status"] != "ISSUED"
        or str(uuid.UUID(statement["report_id"])) != statement["report_id"]
    ):
        raise ValueError("invalid issuance identity")
    for field in ("markdown_sha256", "snapshot_sha256"):
        if not re.fullmatch(r"[0-9a-f]{64}", statement[field]):
            raise ValueError("invalid artifact digest")
    Ed25519PublicKey.from_public_bytes(public_key).verify(
        base64.b64decode(bundle["signature"], validate=True), canonical_bytes(statement)
    )
    for name, field in (
        ("report.md", "markdown_sha256"),
        ("snapshot.json", "snapshot_sha256"),
    ):
        raw = bounded_read(directory / name, 1_000_000)
        if hashlib.sha256(raw).hexdigest() != statement[field]:
            raise ValueError("artifact bytes differ from the signed digest")
    return statement


def _node(left: bytes, right: bytes) -> bytes:
    return hashlib.sha256(b"\x01" + left + right).digest()


def verify_inclusion(
    leaf: bytes, index: int, size: int, hashes_hex: list, root: bytes
) -> bool:
    """RFC 9162 section 2.1.3.2."""
    if index < 0 or index >= size:
        return False
    fn, sn, running = index, size - 1, leaf
    for step_hex in hashes_hex:
        if sn == 0:
            return False
        step = bytes.fromhex(step_hex)
        if fn & 1 or fn == sn:
            running = _node(step, running)
            if not fn & 1:
                while fn and not fn & 1:
                    fn >>= 1
                    sn >>= 1
        else:
            running = _node(running, step)
        fn >>= 1
        sn >>= 1
    return sn == 0 and running == root


def verify_transparency(
    path: Path, statement: dict, log_signing_public_key: str
) -> dict:
    """The entry must name this statement, be signed over it by the trusted
    log-signing companion (SubjectPublicKeyInfo DER, base64), be included in
    a checkpoint and the checkpoint must be signed by the pinned log key.
    Raises ValueError or InvalidSignature."""
    entry = parse(bounded_read(path, 65536))
    if entry.get("schema") != "checkii.transparency-entry.v1":
        raise ValueError("unknown transparency entry schema")
    digest = hashlib.sha256(canonical_bytes(statement)).hexdigest()
    if entry["statement_sha256"] != digest:
        raise ValueError("transparency entry names another statement")
    body = base64.b64decode(entry["body"], validate=True)
    logged = json.loads(body)
    spec = logged["spec"]
    if logged["kind"] != "hashedrekord" or spec["data"]["hash"] != {
        "algorithm": "sha256",
        "value": digest,
    }:
        raise ValueError("logged hash is not this statement's digest")
    pem = base64.b64decode(spec["signature"]["publicKey"]["content"], validate=True)
    companion = serialization.load_pem_public_key(pem)
    logged_key = companion.public_bytes(
        serialization.Encoding.DER, serialization.PublicFormat.SubjectPublicKeyInfo
    )
    if logged_key != base64.b64decode(log_signing_public_key, validate=True):
        raise ValueError("logged public key is not the trusted log-signing companion")
    if not isinstance(companion, ec.EllipticCurvePublicKey) or not isinstance(
        companion.curve, ec.SECP256R1
    ):
        raise ValueError("log-signing companion is not an ECDSA P-256 key")
    companion.verify(
        base64.b64decode(spec["signature"]["content"], validate=True),
        canonical_bytes(statement),
        ec.ECDSA(hashes.SHA256()),
    )
    if entry.get("log_signature") not in (None, spec["signature"]["content"]):
        raise ValueError("log_signature differs from the entry body")
    if entry.get("log_signing_key_id") not in (
        None,
        hashlib.sha256(logged_key).hexdigest(),
    ):
        raise ValueError("log_signing_key_id is not the companion's identifier")
    leaf = hashlib.sha256(b"\x00" + body).digest()
    if not str(entry["entry_uuid"]).endswith(leaf.hex()):
        raise ValueError("entry identifier is not the leaf hash")
    proof = entry["inclusion_proof"]
    root = bytes.fromhex(proof["rootHash"])
    if not verify_inclusion(
        leaf, int(proof["logIndex"]), int(proof["treeSize"]), proof["hashes"], root
    ):
        raise ValueError("inclusion proof does not reach the root")
    note, _, trailer = str(entry["checkpoint"]).partition("\n\n")
    note += "\n"
    lines = note.split("\n")
    origin = lines[0].split(" - ", 1)[0]
    if (
        int(lines[1]) != int(proof["treeSize"])
        or base64.b64decode(lines[2], validate=True) != root
    ):
        raise ValueError("checkpoint names another tree")
    log = KNOWN_LOGS.get(str(entry["log_id"]))
    if log is None or log["origin"] != origin or log["log_url"] != entry.get("log_url"):
        raise ValueError("unknown or mismatched transparency log")
    der = base64.b64decode(log["public_key"], validate=True)
    if hashlib.sha256(der).hexdigest() != entry["log_id"]:
        raise ValueError("pinned log key does not hash to the log id")
    key = serialization.load_pem_public_key(
        b"-----BEGIN PUBLIC KEY-----\n"
        + base64.b64encode(der)
        + b"\n-----END PUBLIC KEY-----\n"
    )
    hint = hashlib.sha256(der).digest()[:4]
    for line in trailer.split("\n"):
        if not line.startswith("\u2014 "):
            continue
        _, name, blob_b64 = line.split(" ", 2)
        blob = base64.b64decode(blob_b64, validate=True)
        if name != origin or blob[:4] != hint:
            continue
        key.verify(blob[4:], note.encode(), ec.ECDSA(hashes.SHA256()))
        return {
            "entry_uuid": entry["entry_uuid"],
            "integrated_time": int(entry["integrated_time"]),
            "origin": origin,
        }
    raise ValueError("checkpoint is not signed by the log key")


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--directory", type=Path, default=Path(__file__).resolve().parent
    )
    parser.add_argument("--trusted-key-id", required=True)
    parser.add_argument("--trusted-public-key", required=True)
    parser.add_argument(
        "--trusted-log-signing-public-key",
        help="SubjectPublicKeyInfo DER, base64, of the key's published "
        "log-signing companion; required when transparency.json is present",
    )
    args = parser.parse_args(argv)
    try:
        statement = verify(args.directory, args.trusted_key_id, args.trusted_public_key)
    except (OSError, ValueError, TypeError, KeyError, RecursionError, InvalidSignature):
        print(
            "INVALID: malformed bundle, untrusted key, invalid signature or changed artifact bytes"
        )
        return 1
    transparency = args.directory / "transparency.json"
    logged = None
    if transparency.exists():
        if not args.trusted_log_signing_public_key:
            print(
                "INVALID: transparency.json is present but no trusted log-signing "
                "companion key was supplied (--trusted-log-signing-public-key)"
            )
            return 1
        try:
            logged = verify_transparency(
                transparency, statement, args.trusted_log_signing_public_key
            )
        except (
            OSError,
            ValueError,
            TypeError,
            KeyError,
            RecursionError,
            InvalidSignature,
        ):
            print("INVALID: transparency.json does not verify for this statement")
            return 1
    print("VALID: exact artifact bytes signed by the supplied trusted key")
    print(f"  report id   {statement['report_id']}")
    print(f"  key id      {args.trusted_key_id}")
    if logged is None:
        print(
            "  log         no transparency.json supplied; nothing checked against the log"
        )
    else:
        print(
            f"  log         {logged['origin']} entry {logged['entry_uuid'][:16]}... at {logged['integrated_time']}"
        )
    print(
        "This synthetic sample proves no scan execution, findings or personal identity."
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())