#!/usr/bin/env python3
"""
blackbox-verify: check that a Servury Blackbox VM is a genuine AMD SEV-SNP
guest running exactly the published measured release, and that the SSH host
key it presents is the one bound into its attestation report.

Nothing here trusts the hosting provider. The report is signed by the CPU
(VCEK, chained to AMD's root), the expected measurement is recomputed locally
from the published artifacts, and the host key is read off the network by
your own ssh-keyscan.

    blackbox-verify --host 203.0.113.5 --release ./release-2026.09
    ssh root@203.0.113.5 attest | blackbox-verify --release ./release-2026.09

Dependencies: python3, cryptography, sev-snp-measure (pip), ssh-keyscan.
"""
import argparse, base64, hashlib, json, os, ssl, struct, subprocess, sys, urllib.request, warnings

try:
    from cryptography.utils import CryptographyDeprecationWarning
    warnings.filterwarnings("ignore", category=CryptographyDeprecationWarning)   # AMD VCEKs carry a zero serial
except Exception:
    pass

KDS = "https://kdsintf.amd.com/vcek/v1/Milan"
EXPECTED_POLICY = 0x30000            # SMT allowed, reserved bit 17, no debug, no migration agent
VCPU_TYPE = "EPYC-v4"
GUEST_FEATURES = 0x1

def die(msg):
    print("FAIL: " + msg)
    sys.exit(1)

def ok(msg):
    print("ok:   " + msg)

def parse_report(r):
    if len(r) < 0x4A0:
        die("report is %d bytes, expected at least 1184" % len(r))
    f = {}
    f["version"], f["guest_svn"], f["policy"] = struct.unpack_from("<IIQ", r, 0)
    f["vmpl"], f["sig_algo"] = struct.unpack_from("<II", r, 0x30)
    f["platform_version"], f["platform_info"], f["flags"] = struct.unpack_from("<QQI", r, 0x38)
    f["report_data"] = r[0x50:0x90]
    f["measurement"] = r[0x90:0xC0]
    f["host_data"] = r[0xC0:0xE0]
    f["id_key_digest"] = r[0xE0:0x110]
    f["author_key_digest"] = r[0x110:0x140]
    f["report_id"] = r[0x140:0x160]
    f["reported_tcb"] = r[0x180:0x188]
    f["chip_id"] = r[0x1A0:0x1E0]
    f["committed_tcb"] = r[0x1E0:0x1E8]
    f["launch_tcb"] = r[0x218:0x220]
    f["signed"] = r[:0x2A0]
    f["sig_r"] = r[0x2A0:0x2A0 + 72]
    f["sig_s"] = r[0x2A0 + 72:0x2A0 + 144]
    return f

def fetch(url):
    ctx = ssl.create_default_context()
    with urllib.request.urlopen(url, timeout=60, context=ctx) as u:
        return u.read()

def verify_signature(f):
    from cryptography import x509
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.asymmetric import ec
    from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature
    tcb = f["reported_tcb"]
    url = "%s/%s?blSPL=%d&teeSPL=%d&snpSPL=%d&ucodeSPL=%d" % (KDS, f["chip_id"].hex(), tcb[0], tcb[1], tcb[6], tcb[7])
    cache = os.path.join(os.path.expanduser("~"), ".cache", "blackbox-verify")
    os.makedirs(cache, exist_ok=True)
    vpath = os.path.join(cache, hashlib.sha256(url.encode()).hexdigest() + ".der")
    if os.path.exists(vpath):
        vcek_der = open(vpath, "rb").read()
    else:
        vcek_der = fetch(url)
        open(vpath, "wb").write(vcek_der)
    chain_path = os.path.join(cache, "milan_cert_chain.pem")
    if os.path.exists(chain_path):
        chain = open(chain_path, "rb").read()
    else:
        chain = fetch(KDS + "/cert_chain")
        open(chain_path, "wb").write(chain)
    vcek = x509.load_der_x509_certificate(vcek_der)
    pems = [p for p in chain.split(b"-----END CERTIFICATE-----") if b"BEGIN CERTIFICATE" in p]
    certs = [x509.load_pem_x509_certificate(p + b"-----END CERTIFICATE-----\n") for p in pems]
    if len(certs) != 2:
        die("AMD cert chain did not contain exactly ASK and ARK")
    ask, ark = certs[0], certs[1]
    # ARK is self-signed; ASK signed by ARK; VCEK signed by ASK. AMD uses RSA-PSS SHA-384.
    from cryptography.hazmat.primitives.asymmetric import padding
    def rsa_verify(cert, signer):
        signer.public_key().verify(cert.signature, cert.tbs_certificate_bytes,
                                   padding.PSS(mgf=padding.MGF1(hashes.SHA384()), salt_length=48), hashes.SHA384())
    try:
        rsa_verify(ark, ark); rsa_verify(ask, ark); rsa_verify(vcek, ask)
    except Exception as e:
        die("certificate chain does not verify: %s" % e)
    ok("VCEK -> ASK -> ARK chain verifies (ARK fingerprint %s)" % hashlib.sha256(ark.public_bytes(serialization.Encoding.DER)).hexdigest()[:32])
    r = int.from_bytes(f["sig_r"][:48], "little")
    s = int.from_bytes(f["sig_s"][:48], "little")
    try:
        vcek.public_key().verify(encode_dss_signature(r, s), f["signed"], ec.ECDSA(hashes.SHA384()))
    except Exception as e:
        die("report signature does not verify with the VCEK: %s" % e)
    ok("report signature verifies (ECDSA P-384, signed inside the CPU)")

def expected_measurement(release, vcpus, cmdline):
    try:
        from sevsnpmeasure import guest, vcpu_types
        from sevsnpmeasure.sev_mode import SevMode
        from sevsnpmeasure.vmm_types import VMMType
    except ImportError:
        die("python module sevsnpmeasure missing: pip install sev-snp-measure")
    ovmf = os.path.join(release, "OVMF.amdsev.fd")
    kernel = os.path.join(release, "vmlinuz")
    initrd = os.path.join(release, "initrd.img")
    for p in (ovmf, kernel, initrd):
        if not os.path.isfile(p):
            die("release artifact missing: " + p)
    ld = guest.calc_launch_digest(SevMode.SEV_SNP, vcpus, vcpu_types.CPU_SIGS[VCPU_TYPE], ovmf, kernel, initrd, cmdline,
                                  GUEST_FEATURES, "", VMMType.QEMU, dump_vmsa=False)
    return ld.hex()

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--host", help="VM address: fetch the attestation over SSH and keyscan the host key yourself")
    ap.add_argument("--port", type=int, default=22)
    ap.add_argument("--release", required=True, help="directory holding OVMF.amdsev.fd, vmlinuz, initrd.img of the release the VM claims")
    ap.add_argument("--vcpus", type=int, help="vCPU count of the VM (default: taken from the panel's attestation JSON if given, else required)")
    ap.add_argument("--panel-json", help="attestation JSON downloaded from the Servury panel (only used for vcpus/expected cross-check)")
    ap.add_argument("--no-kds", action="store_true", help="skip signature verification (offline); NOT a real verification")
    a = ap.parse_args()

    if a.host:
        cmd = ["ssh", "-p", str(a.port), "-o", "BatchMode=yes", "-i", os.environ.get("BB_SSH_KEY", os.path.expanduser("~/.ssh/id_ed25519")), "-o", "StrictHostKeyChecking=no",
               "-o", "UserKnownHostsFile=/dev/null", "-o", "LogLevel=ERROR", "root@" + a.host, "attest"]
        out = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        if out.returncode != 0 or not out.stdout.strip():
            die("could not fetch attestation over SSH: " + out.stderr.strip())
        att = json.loads(out.stdout)
        scan = subprocess.run(["ssh-keyscan", "-p", str(a.port), "-t", "ed25519", a.host], capture_output=True, text=True, timeout=30).stdout
        seen = [l.split()[2] for l in scan.splitlines() if " ssh-ed25519 " in l]
        if not seen:
            die("ssh-keyscan returned no ed25519 host key for %s" % a.host)
        seen_fp = "SHA256:" + base64.b64encode(hashlib.sha256(base64.b64decode(seen[0])).digest()).decode().rstrip("=")
    else:
        att = json.load(sys.stdin)
        seen_fp = None

    if att.get("status") != "ok":
        die("VM reports status %r: it did not produce an SNP report. Do not enter a passphrase." % att.get("status"))
    report = base64.b64decode(att["report_b64"])
    f = parse_report(report)
    print("report: version %d, policy 0x%x, vmpl %d, chip %s..." % (f["version"], f["policy"], f["vmpl"], f["chip_id"].hex()[:16]))

    if f["policy"] != EXPECTED_POLICY:
        die("guest policy is 0x%x, expected 0x%x (debug or migration bits set?)" % (f["policy"], EXPECTED_POLICY))
    ok("policy 0x%x: debugging disabled, no migration agent" % f["policy"])
    if f["vmpl"] != 0:
        die("report was requested from VMPL %d" % f["vmpl"])

    hostkey_blob = base64.b64decode(att["hostkey"].split()[1])
    bound = f["report_data"][:32]
    if hashlib.sha256(hostkey_blob).digest() != bound or f["report_data"][32:] != b"\0" * 32:
        die("REPORT_DATA is not the SHA-256 of the host key the VM sent")
    bound_fp = "SHA256:" + base64.b64encode(bound).decode().rstrip("=")
    ok("REPORT_DATA binds host key %s" % bound_fp)
    if seen_fp is not None:
        if seen_fp != bound_fp:
            die("the host key on the network (%s) is NOT the one in the report (%s): someone is in the middle" % (seen_fp, bound_fp))
        ok("the host key %s presents on the network is that key" % a.host)
    else:
        print("      compare it yourself with the fingerprint your ssh client showed")

    vcpus = a.vcpus
    panel = None
    if a.panel_json:
        panel = json.load(open(a.panel_json))
        vcpus = vcpus or int(panel.get("vcpus", 0))
    if not vcpus:
        die("--vcpus is required (the number of vCPUs the VM was started with)")
    # The firmware's kernel loader appends "initrd=initrd" so the EFI stub can
    # find the initrd it was handed; the hash in the measurement covers the
    # command line as QEMU passed it, without that suffix.
    cmdline = att["cmdline"]
    if cmdline.endswith(" initrd=initrd"):
        cmdline = cmdline[:-len(" initrd=initrd")]
    exp = expected_measurement(a.release, vcpus, cmdline)
    got = f["measurement"].hex()
    if exp != got:
        die("measurement mismatch\n      expected %s\n      report   %s\n      The VM is not running the published release with this command line." % (exp, got))
    ok("launch measurement matches the published release for %d vCPU(s)" % vcpus)
    if panel and panel.get("expected_measurement") and panel["expected_measurement"] != exp:
        die("the panel's expected_measurement differs from your local computation")

    if a.no_kds:
        print("WARN: signature NOT verified (--no-kds); this run proves nothing about the CPU")
    else:
        verify_signature(f)

    print("PASS: %s is a genuine SEV-SNP guest running the measured release; the host key above is safe to trust for this boot." % (a.host or "this VM"))
    if a.host:
        print("      known_hosts line: [%s]:%d ssh-ed25519 %s" % (a.host, a.port, seen[0]) if a.port != 22 else "      known_hosts line: %s ssh-ed25519 %s" % (a.host, seen[0]))

if __name__ == "__main__":
    main()
