vitor-aignosi
0e445de386
QTZPOC-20: Implement .sientia model bundle import workflow Introduces a new Temporal workflow and activities to import .sientia model bundles. This enables: - Secure decryption and validation of bundles through a seven-gate process. - Provisioning imported models in MLflow and updating the MongoDB model listing. - Comprehensive status tracking in the public.experiment_run table with new IMPORT_SUCCESS/IMPORT_ERROR states. - Dedicated metrics, robust error handling, and configurable settings. - Includes updated e2e database fixtures and extensive unit/integration tests for bundle format, reader, and overall workflow.
217 lines
8.9 KiB
Python
217 lines
8.9 KiB
Python
"""The `.sientia` wire format, as the producer writes it.
|
|
|
|
Every constant here is copied from the producer that emits the bundles this reader opens — the
|
|
Streamlit app in `sientia-projects-templates`, modules
|
|
`app/src/operations/model_export/crypto.py` and `app/src/operations/model_export/bundle.py`. The
|
|
producer file each value came from is named next to it. The golden fixtures under
|
|
`tests/fixtures/bundle/` are the anti-drift device: if the producer's layout moves, the tests here
|
|
fail against the committed bytes instead of failing in production.
|
|
|
|
This module deliberately imports nothing from Laborious, Temporal, MLflow or libsodium: it is pure
|
|
`struct` and `hashlib` over bytes, so it stays liftable into a shared package if the format ever
|
|
gets one (the cancelled QTZPOC-15). `reader.py` is where the crypto lives.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import struct
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
# --- Producer: app/src/operations/model_export/crypto.py ---
|
|
MAGIC: Final[bytes] = b'SIENTIA1'
|
|
FORMAT_VERSION: Final[int] = 1
|
|
AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM: Final[int] = 1
|
|
KDF_ID_ARGON2ID: Final[int] = 1
|
|
SALT_SIZE: Final[int] = 16
|
|
DIGEST_SIZE: Final[int] = 32
|
|
RESERVED_SIGNATURE_SIZE: Final[int] = 64
|
|
# `KDF_MEMLIMIT_BYTES` / `KDF_OPSLIMIT` on the producer: what it writes is also the ceiling the
|
|
# reader refuses to exceed, so header bytes cannot buy an arbitrary allocation on the worker.
|
|
MAX_KDF_MEMLIMIT: Final[int] = 256 * 1024 * 1024 # 268435456
|
|
MAX_KDF_OPSLIMIT: Final[int] = 3
|
|
# `KDF_PARALLELISM_RESERVED`: libsodium's `crypto_pwhash` has no lane parameter, so the field is
|
|
# reserved and must be exactly 1.
|
|
PARALLELISM: Final[int] = 1
|
|
CHUNK_SIZE: Final[int] = 64 * 1024 # 65536, fixed by the format
|
|
KEY_SIZE: Final[int] = 32 # crypto_secretstream_xchacha20poly1305_KEYBYTES
|
|
|
|
# magic(8s) version(B) aead_id(B) kdf_id(B) salt(16s) memlimit(Q) opslimit(Q)
|
|
# parallelism(B) digest(32s) reserved_signature(64s) — producer's `_HEADER_STRUCT`.
|
|
HEADER_STRUCT: Final[struct.Struct] = struct.Struct('>8sBBB16sQQB32s64s')
|
|
HEADER_SIZE: Final[int] = HEADER_STRUCT.size # 140
|
|
|
|
# --- libsodium constants the framing depends on ---
|
|
# Not producer constants, but constants of the primitive the producer uses. Kept as literals so
|
|
# this module stays import-free; `tests/laborious/utils/bundle/test_format.py` asserts they equal
|
|
# the values `nacl.bindings` reports, which is what makes the literals safe.
|
|
STREAM_HEADER_SIZE: Final[int] = 24 # crypto_secretstream_xchacha20poly1305_HEADERBYTES
|
|
ABYTES: Final[int] = 17 # crypto_secretstream_xchacha20poly1305_ABYTES
|
|
|
|
# The pre-upload check prefix of `08-security-and-encryption.md` § 5.3: enough bytes to hold the
|
|
# header, the stream header and one full chunk. Derived from constants, never read from the file.
|
|
PREFIX_SIZE: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE + CHUNK_SIZE + ABYTES # 65717
|
|
|
|
# The digest covers the chunk stream only. The producer computes `sha256(ciphertext)` over the
|
|
# joined chunks and *then* writes `header + stream_header + ciphertext`, so on disk the covered
|
|
# range starts after both headers. Verified against a bundle produced by the producer's own code:
|
|
# `sha256(file[164:])` matches the header field and `sha256(file[140:])` does not.
|
|
DIGEST_OFFSET: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE # 164
|
|
|
|
MIN_FILE_SIZE: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE
|
|
|
|
SUPPORTED_FORMAT_VERSIONS: Final[frozenset[int]] = frozenset({FORMAT_VERSION})
|
|
SUPPORTED_AEAD_IDS: Final[frozenset[int]] = frozenset({AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM})
|
|
SUPPORTED_KDF_IDS: Final[frozenset[int]] = frozenset({KDF_ID_ARGON2ID})
|
|
|
|
_HASH_BLOCK_SIZE: Final[int] = 1024 * 1024
|
|
|
|
|
|
class BundleFormatError(Exception):
|
|
"""The bytes are not a `.sientia` bundle this reader implements.
|
|
|
|
Raised for every header-level rejection: wrong magic, unknown version, unknown primitive ids,
|
|
cost parameters above the ceiling, a reserved field carrying a value, a file too short to hold
|
|
a header, or a chunk stream whose digest does not match the header's. `reader.py` turns this
|
|
into a gate-2 rejection; nothing here knows about gates or import steps.
|
|
"""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BundleHeader:
|
|
"""The 140-byte plaintext header, parsed and validated."""
|
|
|
|
format_version: int
|
|
aead_id: int
|
|
kdf_id: int
|
|
salt: bytes
|
|
kdf_memlimit_bytes: int
|
|
kdf_opslimit: int
|
|
kdf_parallelism: int
|
|
ciphertext_digest: bytes
|
|
reserved_signature: bytes
|
|
|
|
|
|
def parse_header(raw: bytes) -> BundleHeader:
|
|
"""Parse and validate the plaintext header.
|
|
|
|
Everything the reader trusts about the file's shape is decided here, before any offset past
|
|
the header is used and before a key is derived.
|
|
|
|
Args:
|
|
raw: at least `HEADER_SIZE` bytes read from the start of the object.
|
|
|
|
Returns:
|
|
BundleHeader: the validated header.
|
|
|
|
Raises:
|
|
BundleFormatError: on any invalid or unsupported field.
|
|
"""
|
|
if len(raw) < HEADER_SIZE:
|
|
raise BundleFormatError(
|
|
f'file is shorter than the {HEADER_SIZE}-byte header ({len(raw)} bytes available)'
|
|
)
|
|
|
|
(
|
|
magic,
|
|
format_version,
|
|
aead_id,
|
|
kdf_id,
|
|
salt,
|
|
kdf_memlimit_bytes,
|
|
kdf_opslimit,
|
|
kdf_parallelism,
|
|
ciphertext_digest,
|
|
reserved_signature,
|
|
) = HEADER_STRUCT.unpack(raw[:HEADER_SIZE])
|
|
|
|
if magic != MAGIC:
|
|
raise BundleFormatError('file does not start with the expected bundle marker')
|
|
if format_version not in SUPPORTED_FORMAT_VERSIONS:
|
|
raise BundleFormatError(f'unsupported bundle format version: {format_version}')
|
|
if aead_id not in SUPPORTED_AEAD_IDS:
|
|
raise BundleFormatError(f'unsupported encryption identifier: {aead_id}')
|
|
if kdf_id not in SUPPORTED_KDF_IDS:
|
|
raise BundleFormatError(f'unsupported key derivation identifier: {kdf_id}')
|
|
if kdf_memlimit_bytes > MAX_KDF_MEMLIMIT:
|
|
raise BundleFormatError(
|
|
f'declared key derivation memory {kdf_memlimit_bytes} exceeds the '
|
|
f'{MAX_KDF_MEMLIMIT} ceiling'
|
|
)
|
|
if kdf_opslimit > MAX_KDF_OPSLIMIT:
|
|
raise BundleFormatError(
|
|
f'declared key derivation passes {kdf_opslimit} exceeds the {MAX_KDF_OPSLIMIT} ceiling'
|
|
)
|
|
if kdf_memlimit_bytes <= 0 or kdf_opslimit <= 0:
|
|
raise BundleFormatError('key derivation parameters must be positive')
|
|
if kdf_parallelism != PARALLELISM:
|
|
raise BundleFormatError(
|
|
f'reserved parallelism field must be {PARALLELISM}, found {kdf_parallelism}'
|
|
)
|
|
if reserved_signature != bytes(RESERVED_SIGNATURE_SIZE):
|
|
raise BundleFormatError('reserved signature field must be empty in this format version')
|
|
|
|
return BundleHeader(
|
|
format_version=format_version,
|
|
aead_id=aead_id,
|
|
kdf_id=kdf_id,
|
|
salt=salt,
|
|
kdf_memlimit_bytes=kdf_memlimit_bytes,
|
|
kdf_opslimit=kdf_opslimit,
|
|
kdf_parallelism=kdf_parallelism,
|
|
ciphertext_digest=ciphertext_digest,
|
|
reserved_signature=reserved_signature,
|
|
)
|
|
|
|
|
|
def read_header(path: Path | str) -> BundleHeader:
|
|
"""Read and validate the header of the file at `path`.
|
|
|
|
Raises:
|
|
BundleFormatError: when the file cannot hold a header plus a stream header, or when the
|
|
header itself is invalid.
|
|
"""
|
|
file_path = Path(path)
|
|
size = file_path.stat().st_size
|
|
if size < MIN_FILE_SIZE:
|
|
raise BundleFormatError(
|
|
f'file is {size} bytes, smaller than the {MIN_FILE_SIZE} bytes a bundle needs for its '
|
|
'header and encrypted stream header'
|
|
)
|
|
with file_path.open('rb') as handle:
|
|
return parse_header(handle.read(HEADER_SIZE))
|
|
|
|
|
|
def compute_ciphertext_digest(path: Path | str) -> bytes:
|
|
"""SHA-256 of the chunk stream — the bytes from `DIGEST_OFFSET` to the end of the file."""
|
|
digest = hashlib.sha256()
|
|
with Path(path).open('rb') as handle:
|
|
handle.seek(DIGEST_OFFSET)
|
|
while block := handle.read(_HASH_BLOCK_SIZE):
|
|
digest.update(block)
|
|
return digest.digest()
|
|
|
|
|
|
def verify_ciphertext_digest(path: Path | str, header: BundleHeader) -> None:
|
|
"""Compare the file's chunk-stream digest against the header's.
|
|
|
|
Raises:
|
|
BundleFormatError: when they differ — a truncated, overwritten or corrupted transfer. This
|
|
runs before any key derivation, so a damaged file never costs an Argon2id pass.
|
|
"""
|
|
if compute_ciphertext_digest(path) != header.ciphertext_digest:
|
|
raise BundleFormatError(
|
|
'the encrypted content does not match the digest recorded in the bundle header'
|
|
)
|
|
|
|
|
|
def file_digest(path: Path | str) -> str:
|
|
"""Hex SHA-256 of the whole file, for the `expected_digest` comparison of gate 1."""
|
|
digest = hashlib.sha256()
|
|
with Path(path).open('rb') as handle:
|
|
while block := handle.read(_HASH_BLOCK_SIZE):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|