"""Check the two raw NRAS tokens of this folder OUTSIDE the VM, with NVIDIA's own keys.

For each token: the ES384 signature of the summary JWT and of every per-device JWT against the
key NVIDIA publishes at https://nras.attestation.nvidia.com/.well-known/jwks.json (fetched live),
the issuer, the overall result, eat_nonce equal to the nonce drawn outside the VM, and per device
measres success, secure boot on, debug disabled.

    pip install "PyJWT[crypto]>=2.7"
    python check_raw_tokens.py nras_token.json <gpu nonce> nras_switch_token.json <switch nonce>

The nonces are printed in attestation-run.txt and nvswitch-diag.txt ("Tenant nonce, drawn outside the VM").
"""
import json
import sys
import urllib.request

import jwt
from jwt import PyJWKSet

JWKS_URL = "https://nras.attestation.nvidia.com/.well-known/jwks.json"
ISSUER = "https://nras.attestation.nvidia.com"
SHOWN = (
    "hwmodel", "measres", "secboot", "dbgstat",
    "x-nvidia-gpu-driver-version", "x-nvidia-gpu-vbios-version", "x-nvidia-switch-bios-version",
    "x-nvidia-gpu-attestation-report-nonce-match", "x-nvidia-switch-attestation-report-nonce-match",
    "x-nvidia-gpu-attestation-report-signature-verified", "x-nvidia-switch-attestation-report-signature-verified",
    "x-nvidia-gpu-attestation-report-cert-chain-validated", "x-nvidia-switch-attestation-report-cert-chain-validated",
)


def fetch_keys() -> dict:
    req = urllib.request.Request(JWKS_URL, headers={"User-Agent": "evidence-check/1"})
    with urllib.request.urlopen(req, timeout=30) as resp:
        jwks = json.loads(resp.read().decode("utf-8"))
    return {k.key_id: k for k in PyJWKSet.from_dict(jwks).keys}


def decode(token: str, keys: dict) -> dict:
    kid = jwt.get_unverified_header(token)["kid"]
    if kid not in keys:
        raise ValueError(f"kid {kid} is not in NVIDIA's JWKS")
    return jwt.decode(token, key=keys[kid].key, algorithms=["ES384"],
                      options={"verify_aud": False, "verify_exp": False, "verify_nbf": False})


def check(path: str, nonce: str, keys: dict, label: str) -> bool:
    data = json.load(open(path, encoding="utf-8"))
    # GPU file (SDK token): [["JWT", local envelope], {"REMOTE_GPU_CLAIMS": [["JWT", summary], {"GPU-0": jwt, ...}]}]
    # Switch file (raw NRAS answer): [["JWT", summary], {"SWITCH-0": jwt, ...}]
    nras = data[1]["REMOTE_GPU_CLAIMS"] if isinstance(data[1], dict) and "REMOTE_GPU_CLAIMS" in data[1] else data
    summary_jwt, devices = nras[0][1], nras[1]
    s = decode(summary_jwt, keys)
    ok = s.get("iss") == ISSUER and s.get("x-nvidia-overall-att-result") is True and s.get("eat_nonce") == nonce
    print(f"[{label}] summary: ES384 signature OK (kid {jwt.get_unverified_header(summary_jwt)['kid']}), "
          f"iss {s.get('iss')}, overall {s.get('x-nvidia-overall-att-result')}, eat_nonce matches: {s.get('eat_nonce') == nonce}")
    for name in sorted(devices):
        p = decode(devices[name], keys)
        dev_ok = (p.get("eat_nonce") == nonce and p.get("measres") == "success"
                  and p.get("secboot") is True and p.get("dbgstat") == "disabled")
        ok = ok and dev_ok
        shown = {k.replace("x-nvidia-", ""): p[k] for k in SHOWN if k in p}
        print(f"[{label}] {name}: ES384 signature OK, eat_nonce matches: {p.get('eat_nonce') == nonce}, {json.dumps(shown, sort_keys=True)}")
    print(f"[{label}] RESULT: {'ALL CHECKS PASS' if ok else 'FAIL'}")
    return ok


def main() -> int:
    keys = fetch_keys()
    print(f"NVIDIA JWKS {JWKS_URL}: {len(keys)} keys")
    gpu = check(sys.argv[1], sys.argv[2], keys, "GPU")
    switch = check(sys.argv[3], sys.argv[4], keys, "NVSwitch")
    return 0 if gpu and switch else 1


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