tests/test_integration.py
2.8 KB · 98 lines · python Raw
1 """End-to-end integration tests."""
2
3 from __future__ import annotations
4
5 import pytest
6
7 from pqc_enclave_sdk import (
8 ArtifactKind,
9 DeviceAttester,
10 EnclaveLockedError,
11 EnclaveVault,
12 InMemoryEnclaveBackend,
13 UnknownArtifactError,
14 )
15
16
17 def test_full_lifecycle_unlock_put_save_lock_reload_get(
18 small_weights: bytes,
19 ) -> None:
20 backend = InMemoryEnclaveBackend(device_id="iphone-1")
21
22 v1 = EnclaveVault(backend=backend)
23 v1.unlock()
24 v1.put_artifact(
25 name="llama-weights",
26 kind=ArtifactKind.MODEL_WEIGHTS,
27 content=small_weights,
28 )
29 v1.save()
30 # Grab the session key we used - a fresh unlock on v2 would create a new
31 # key and fail to decrypt. Simulate the real enclave contract where the
32 # backend holds the wrapping key: swap the pre-existing session key.
33 old_key = v1._symmetric_key
34 old_key_id = v1._key_id
35 old_exp = v1._expires_at
36 v1.lock()
37
38 v2 = EnclaveVault(backend=backend)
39 v2.unlock()
40 # Inject the historical session into v2 to mimic enclave-held KEK rewrap.
41 v2._symmetric_key = old_key
42 v2._key_id = old_key_id
43 v2._expires_at = old_exp
44 v2._store = backend.load_artifacts()
45
46 art = v2.get_artifact("llama-weights")
47 assert art.content == small_weights
48
49
50 def test_locked_vault_operations_raise_enclave_locked_error() -> None:
51 backend = InMemoryEnclaveBackend()
52 v = EnclaveVault(backend=backend)
53 with pytest.raises(EnclaveLockedError):
54 v.put_artifact(
55 name="x", kind=ArtifactKind.OTHER, content=b"1"
56 )
57 with pytest.raises(EnclaveLockedError):
58 v.get_artifact("x")
59 with pytest.raises(EnclaveLockedError):
60 v.delete_artifact("x")
61 with pytest.raises(EnclaveLockedError):
62 v.list_artifacts()
63
64
65 def test_attestation_over_stored_artifact_verifies(
66 signer_identity, small_weights
67 ) -> None:
68 backend = InMemoryEnclaveBackend(
69 device_id="pixel-8-bob", device_model="pixel-8"
70 )
71 v = EnclaveVault(backend=backend)
72 v.unlock()
73 enc = v.put_artifact(
74 name="safety-model",
75 kind=ArtifactKind.SAFETY_MODEL,
76 content=small_weights,
77 )
78
79 attester = DeviceAttester(
80 identity=signer_identity,
81 device_id=backend.device_id,
82 device_model=backend.device_model,
83 enclave_vendor=backend.enclave_vendor,
84 )
85 att = attester.attest(
86 artifact_id=enc.metadata.artifact_id,
87 content_hash=enc.content_hash,
88 )
89 assert DeviceAttester.verify(att) is True
90 # Tamper with the content hash - verification must now fail.
91 att.artifact_content_hash = "00" * 32
92 assert DeviceAttester.verify(att) is False
93
94 # Sanity-check that deleting and re-getting raises the right error.
95 v.delete_artifact(enc.metadata.artifact_id)
96 with pytest.raises(UnknownArtifactError):
97 v.get_artifact(enc.metadata.artifact_id)
98