src/pqc_enclave_sdk/artifact.py
2.8 KB · 103 lines · python Raw
1 """On-device AI artifacts (model weights, credentials, adapters) - data structures."""
2
3 from __future__ import annotations
4
5 import hashlib
6 from dataclasses import asdict, dataclass
7 from enum import Enum
8 from typing import Any
9
10
11 class ArtifactKind(str, Enum):
12 """Kinds of artifacts stored in a device secure enclave."""
13
14 MODEL_WEIGHTS = "model-weights"
15 LORA_ADAPTER = "lora-adapter"
16 TOKENIZER = "tokenizer"
17 CREDENTIAL = "credential"
18 BIOMETRIC_TEMPLATE = "biometric-template"
19 INFERENCE_CACHE = "inference-cache"
20 SAFETY_MODEL = "safety-model"
21 OTHER = "other"
22
23
24 @dataclass(frozen=True)
25 class ArtifactMetadata:
26 """Non-secret metadata about an artifact."""
27
28 artifact_id: str
29 name: str
30 kind: ArtifactKind
31 version: str = ""
32 app_bundle_id: str = ""
33 size_bytes: int = 0
34 created_at: str = ""
35 device_id: str = ""
36 model_did: str = ""
37 tags: tuple[str, ...] = ()
38 description: str = ""
39
40 def to_dict(self) -> dict[str, Any]:
41 d = asdict(self)
42 d["kind"] = self.kind.value
43 d["tags"] = list(self.tags)
44 return d
45
46
47 @dataclass
48 class EnclaveArtifact:
49 """A plaintext artifact - only exists inside an unlocked EnclaveVault."""
50
51 metadata: ArtifactMetadata
52 content: bytes
53
54 @staticmethod
55 def content_hash(data: bytes) -> str:
56 return hashlib.sha3_256(data).hexdigest()
57
58 def sha3_256_hex(self) -> str:
59 return self.content_hash(self.content)
60
61
62 @dataclass
63 class EncryptedArtifact:
64 """AES-256-GCM encrypted artifact stored in the enclave-backed store."""
65
66 metadata: ArtifactMetadata
67 nonce: str
68 ciphertext: str
69 content_hash: str
70 key_id: str
71
72 def to_dict(self) -> dict[str, Any]:
73 return {
74 "metadata": self.metadata.to_dict(),
75 "nonce": self.nonce,
76 "ciphertext": self.ciphertext,
77 "content_hash": self.content_hash,
78 "key_id": self.key_id,
79 }
80
81 @classmethod
82 def from_dict(cls, data: dict[str, Any]) -> EncryptedArtifact:
83 meta = data["metadata"]
84 return cls(
85 metadata=ArtifactMetadata(
86 artifact_id=meta["artifact_id"],
87 name=meta["name"],
88 kind=ArtifactKind(meta["kind"]),
89 version=meta.get("version", ""),
90 app_bundle_id=meta.get("app_bundle_id", ""),
91 size_bytes=int(meta.get("size_bytes", 0)),
92 created_at=meta.get("created_at", ""),
93 device_id=meta.get("device_id", ""),
94 model_did=meta.get("model_did", ""),
95 tags=tuple(meta.get("tags", [])),
96 description=meta.get("description", ""),
97 ),
98 nonce=data["nonce"],
99 ciphertext=data["ciphertext"],
100 content_hash=data["content_hash"],
101 key_id=data["key_id"],
102 )
103