examples/update_rollback_blocked.py
1.9 KB · 64 lines · python Raw
1 """Example: UpdateChain blocks a rollback from v1.0 to v0.9.
2
3 Rollback attacks abuse legitimately-signed old firmware to re-introduce known
4 vulnerabilities. UpdateChain enforces monotonic versions by default. Set
5 `allow_rollback=True` only for emergency break-glass scenarios.
6 """
7
8 from __future__ import annotations
9
10 from quantumshield.identity.agent import AgentIdentity
11
12 from pqc_bootloader import (
13 FirmwareImage,
14 FirmwareMetadata,
15 FirmwareRollbackError,
16 FirmwareSigner,
17 TargetDevice,
18 UpdateChain,
19 )
20
21
22 def build(name: str, version: str, payload: bytes) -> FirmwareImage:
23 meta = FirmwareMetadata(
24 name=name,
25 version=version,
26 target=TargetDevice.AI_INFERENCE_APPLIANCE,
27 )
28 return FirmwareImage.from_bytes(meta, payload)
29
30
31 def main() -> None:
32 signer = FirmwareSigner(AgentIdentity.create("acme-appliance-vendor"))
33 chain = UpdateChain()
34
35 v1 = signer.sign(build("acme-inference-os", "1.0.0", b"v1.0 payload"))
36 chain.add(v1)
37 print(f"[chain] added v{v1.firmware.metadata.version} (hash={v1.firmware.image_hash[:16]}...)")
38
39 # Attacker tries to push v0.9 as the next update (rollback).
40 v0 = signer.sign(
41 build("acme-inference-os", "0.9.0", b"v0.9 payload"),
42 previous_firmware_hash=v1.firmware.image_hash,
43 )
44 try:
45 chain.add(v0)
46 except FirmwareRollbackError as exc:
47 print(f"[chain] BLOCKED v{v0.firmware.metadata.version}: {exc}")
48
49 # Forward update is fine.
50 v2 = signer.sign(
51 build("acme-inference-os", "1.1.0", b"v1.1 payload"),
52 previous_firmware_hash=v1.firmware.image_hash,
53 )
54 chain.add(v2)
55 print(f"[chain] added v{v2.firmware.metadata.version} (hash={v2.firmware.image_hash[:16]}...)")
56
57 ok, errors = chain.verify_chain()
58 print(f"[chain] verify_chain: ok={ok} errors={errors}")
59 print(f"[chain] current version = {chain.current().firmware.metadata.version}")
60
61
62 if __name__ == "__main__":
63 main()
64