examples/driver_attestation.py
| 1 | """Sign and verify a GPU driver module attestation with an allow-list. |
| 2 | |
| 3 | Run: |
| 4 | |
| 5 | python examples/driver_attestation.py |
| 6 | """ |
| 7 | |
| 8 | from __future__ import annotations |
| 9 | |
| 10 | from quantumshield.identity.agent import AgentIdentity |
| 11 | |
| 12 | from pqc_gpu_driver import ( |
| 13 | DriverAttestationVerifier, |
| 14 | DriverAttester, |
| 15 | DriverModule, |
| 16 | ) |
| 17 | |
| 18 | |
| 19 | def main() -> None: |
| 20 | # A fake nvidia.ko blob. |
| 21 | driver_bytes = b"\x7fELF\x02\x01\x01\x00" + b"\x00" * 8 + b"NVIDIA-GPU-DRV" * 256 |
| 22 | module = DriverModule( |
| 23 | name="nvidia.ko", |
| 24 | version="550.54.14", |
| 25 | module_hash=DriverModule.hash_module_bytes(driver_bytes), |
| 26 | module_size=len(driver_bytes), |
| 27 | target="linux", |
| 28 | ) |
| 29 | |
| 30 | vendor = AgentIdentity.create("nvidia-driver-signer", capabilities=["attest"]) |
| 31 | attacker = AgentIdentity.create("rogue-signer", capabilities=["attest"]) |
| 32 | |
| 33 | print("[*] Vendor signing driver module with ML-DSA ...") |
| 34 | attester = DriverAttester(vendor) |
| 35 | attestation = attester.attest(module) |
| 36 | print(f" module = {attestation.module.name} v{attestation.module.version}") |
| 37 | print(f" module_hash = {attestation.module.module_hash[:32]}...") |
| 38 | print(f" signer_did = {attestation.signer_did}") |
| 39 | print(f" algorithm = {attestation.algorithm}") |
| 40 | print(f" signed_at = {attestation.signed_at}") |
| 41 | |
| 42 | verifier = DriverAttestationVerifier(trusted_signers={vendor.did}) |
| 43 | |
| 44 | print("\n[*] Case 1: vendor attestation with correct bytes ...") |
| 45 | result = verifier.verify(attestation, actual_module_bytes=driver_bytes) |
| 46 | print(f" valid = {result.valid}") |
| 47 | print(f" trusted = {result.trusted}") |
| 48 | assert result.valid |
| 49 | |
| 50 | print("\n[*] Case 2: attacker's attestation rejected by allow-list ...") |
| 51 | rogue_att = DriverAttester(attacker).attest(module) |
| 52 | bad = verifier.verify(rogue_att, actual_module_bytes=driver_bytes) |
| 53 | print(f" valid = {bad.valid}") |
| 54 | print(f" error = {bad.error}") |
| 55 | assert not bad.valid |
| 56 | |
| 57 | print("\n[+] Attestation flow verified. Untrusted signers cannot load drivers.") |
| 58 | |
| 59 | |
| 60 | if __name__ == "__main__": |
| 61 | main() |
| 62 | |