QTZPOC-20
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.
This commit is contained in:
272
tests/helpers/bundle_factory.py
Normal file
272
tests/helpers/bundle_factory.py
Normal file
@@ -0,0 +1,272 @@
|
||||
"""Test-only `.sientia` writer, mirrored field for field from the producer.
|
||||
|
||||
There is no production writer in this repository (change `import-model-workflow`, design D1): the
|
||||
producer is the Streamlit app in `sientia-projects-templates`. This factory exists for one reason —
|
||||
building the **hostile** variants the gate tests need, which the producer cannot emit. Its output is
|
||||
pinned against the golden fixture by `test_format.py::test_factory_header_matches_fixture_layout`,
|
||||
so a drift in either direction shows up as a failing test.
|
||||
|
||||
Mirrored from:
|
||||
- `app/src/operations/model_export/crypto.py` — header layout, KDF, `secretstream` framing
|
||||
- `app/src/operations/model_export/bundle.py` — zip layout and the two-top-level-key JSON
|
||||
|
||||
Deliberate deviation: the default KDF cost is a fraction of the producer's (8 MiB / 1 pass instead
|
||||
of 256 MiB / 3), because the reader honours whatever the header declares and a test suite that
|
||||
derives thirty 256 MiB Argon2id keys is a test suite nobody runs. The golden fixture is what covers
|
||||
the production cost parameters.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import struct
|
||||
import unicodedata
|
||||
import zipfile
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import nacl.bindings as sodium
|
||||
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||||
|
||||
# --- mirrored from the producer's crypto.py ---
|
||||
MAGIC = b'SIENTIA1'
|
||||
FORMAT_VERSION = 1
|
||||
AEAD_ID = 1
|
||||
KDF_ID = 1
|
||||
SALT_SIZE = 16
|
||||
RESERVED_SIGNATURE_SIZE = 64
|
||||
CHUNK_SIZE = 64 * 1024
|
||||
HEADER_STRUCT = struct.Struct('>8sBBB16sQQB32s64s')
|
||||
HEADER_SIZE = HEADER_STRUCT.size
|
||||
PRODUCER_KDF_MEMLIMIT_BYTES = 256 * 1024 * 1024
|
||||
PRODUCER_KDF_OPSLIMIT = 3
|
||||
|
||||
# Cheap-but-valid cost parameters for the suite; see the module docstring.
|
||||
TEST_KDF_MEMLIMIT_BYTES = 8 * 1024 * 1024
|
||||
TEST_KDF_OPSLIMIT = 1
|
||||
|
||||
DEFAULT_PASSWORD = 'Factory1Model!Test'
|
||||
DEFAULT_ENVELOPE_KEY = bytes(range(32))
|
||||
|
||||
MODEL_DIRS = ('prediction_model', 'data_model')
|
||||
MODEL_FILES = ('MLmodel', 'model.pkl', 'conda.yaml', 'python_env.yaml', 'requirements.txt')
|
||||
CSV_FILES = ('train_data.csv', 'test_data.csv', 'raw_data.csv')
|
||||
|
||||
DEFAULT_METADATA = {
|
||||
'model_name': 'factory_model',
|
||||
'experiment_name': 'factory_experiment',
|
||||
'run_id': 'aa11bb22cc33dd44ee55ff6677889900',
|
||||
'model_version': '2',
|
||||
'model_project': 'factory_project',
|
||||
'export_timestamp': '2026-08-25T12:00:00+00:00',
|
||||
}
|
||||
|
||||
DEFAULT_PARAMETERS = {
|
||||
'model_name': 'factory_model',
|
||||
'target_variable': 'SE',
|
||||
'input_variables': "['feature_a', 'feature_b']",
|
||||
'retrain': 'False',
|
||||
}
|
||||
|
||||
|
||||
def default_tree() -> dict[str, bytes]:
|
||||
"""The artifact tree a well-formed bundle carries, as archive-name -> bytes."""
|
||||
files: dict[str, bytes] = {}
|
||||
for directory in MODEL_DIRS:
|
||||
for name in MODEL_FILES:
|
||||
files[f'artifacts/{directory}/{name}'] = f'{directory}/{name}\n'.encode()
|
||||
files['artifacts/model_card/model_card.json'] = json.dumps({'target': 'SE'}).encode()
|
||||
for name in CSV_FILES:
|
||||
files[f'artifacts/{name}'] = b'datetime,feature_a,SE\n2026-08-25T00:00:00,1.0,2.0\n'
|
||||
return files
|
||||
|
||||
|
||||
def build_metadata_document(
|
||||
*,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
parameters: dict[str, Any] | None = None,
|
||||
drop_metadata_field: str | None = None,
|
||||
drop_top_level_key: str | None = None,
|
||||
drop_parameter: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build `metadata.json`, with hooks for the malformed variants gate 6 must reject."""
|
||||
meta = {**DEFAULT_METADATA, **(metadata or {})}
|
||||
params = {**DEFAULT_PARAMETERS, **(parameters or {})}
|
||||
if drop_metadata_field:
|
||||
meta.pop(drop_metadata_field, None)
|
||||
if drop_parameter:
|
||||
params.pop(drop_parameter, None)
|
||||
document: dict[str, Any] = {'parameters': params, 'metadata': meta}
|
||||
if drop_top_level_key:
|
||||
document.pop(drop_top_level_key, None)
|
||||
return document
|
||||
|
||||
|
||||
@dataclass
|
||||
class ZipSpec:
|
||||
"""How to lay out the zip inside the bundle."""
|
||||
|
||||
files: dict[str, bytes] = field(default_factory=default_tree)
|
||||
document: dict[str, Any] | None = None
|
||||
include_metadata: bool = True
|
||||
symlinks: dict[str, str] = field(default_factory=dict)
|
||||
absolute_entry: str | None = None
|
||||
traversal_entry: str | None = None
|
||||
bomb_bytes: int = 0
|
||||
extra_files: dict[str, bytes] = field(default_factory=dict)
|
||||
|
||||
|
||||
def build_zip(spec: ZipSpec | None = None) -> bytes:
|
||||
"""Zip `artifacts/` plus `metadata.json`, mirroring the producer's `build_zip`.
|
||||
|
||||
`metadata.json` is written first with `json.dumps` and no indent, entries use `/` separators,
|
||||
and no directory entries are stored — all three matter to the reader.
|
||||
"""
|
||||
spec = spec or ZipSpec()
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, 'w', zipfile.ZIP_DEFLATED) as archive:
|
||||
if spec.include_metadata:
|
||||
document = spec.document if spec.document is not None else build_metadata_document()
|
||||
archive.writestr('metadata.json', json.dumps(document))
|
||||
for name, content in {**spec.files, **spec.extra_files}.items():
|
||||
archive.writestr(_regular_file_info(name), content)
|
||||
for name, target in spec.symlinks.items():
|
||||
info = zipfile.ZipInfo(name)
|
||||
info.create_system = 3
|
||||
info.external_attr = (0o120777 << 16) | 0o200000
|
||||
archive.writestr(info, target)
|
||||
if spec.absolute_entry:
|
||||
archive.writestr(_regular_file_info(spec.absolute_entry), b'absolute\n')
|
||||
if spec.traversal_entry:
|
||||
archive.writestr(_regular_file_info(spec.traversal_entry), b'traversal\n')
|
||||
if spec.bomb_bytes:
|
||||
archive.writestr(_regular_file_info('artifacts/bomb.csv'), b'\0' * spec.bomb_bytes)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def _regular_file_info(name: str) -> zipfile.ZipInfo:
|
||||
"""A zip entry carrying a Unix regular-file mode, as `ZipFile.write` would store it."""
|
||||
info = zipfile.ZipInfo(name)
|
||||
info.create_system = 3
|
||||
info.compress_type = zipfile.ZIP_DEFLATED
|
||||
info.external_attr = 0o100644 << 16
|
||||
return info
|
||||
|
||||
|
||||
def derive_key(password: str, salt: bytes, opslimit: int, memlimit: int) -> bytes:
|
||||
"""Argon2id over the NFC-normalised password — the producer's `_derive_key`."""
|
||||
normalized = unicodedata.normalize('NFC', password)
|
||||
return sodium.crypto_pwhash_alg(
|
||||
sodium.crypto_secretstream_xchacha20poly1305_KEYBYTES,
|
||||
normalized.encode('utf-8'),
|
||||
salt,
|
||||
opslimit,
|
||||
memlimit,
|
||||
sodium.crypto_pwhash_ALG_ARGON2ID13,
|
||||
)
|
||||
|
||||
|
||||
def encrypt_archive( # noqa: PLR0913 - every knob here exists to build one hostile variant
|
||||
plaintext: bytes,
|
||||
password: str = DEFAULT_PASSWORD,
|
||||
*,
|
||||
format_version: int = FORMAT_VERSION,
|
||||
aead_id: int = AEAD_ID,
|
||||
kdf_id: int = KDF_ID,
|
||||
memlimit: int = TEST_KDF_MEMLIMIT_BYTES,
|
||||
opslimit: int = TEST_KDF_OPSLIMIT,
|
||||
declared_memlimit: int | None = None,
|
||||
declared_opslimit: int | None = None,
|
||||
parallelism: int = 1,
|
||||
reserved_signature: bytes | None = None,
|
||||
magic: bytes = MAGIC,
|
||||
digest_override: bytes | None = None,
|
||||
salt: bytes | None = None,
|
||||
drop_final_tag: bool = False,
|
||||
) -> bytes:
|
||||
"""Produce a full `.sientia` byte stream, mirroring the producer's `encrypt_archive`.
|
||||
|
||||
The `declared_*` arguments write cost parameters into the header that differ from the ones
|
||||
actually used, which is how a "512 MiB header" is built without allocating 512 MiB.
|
||||
`drop_final_tag` pushes every chunk as `MESSAGE`, so the stream is framed as if more were
|
||||
coming — the truncation the digest cannot see.
|
||||
"""
|
||||
salt = salt if salt is not None else os.urandom(SALT_SIZE)
|
||||
key = derive_key(password, salt, opslimit, memlimit)
|
||||
|
||||
state = sodium.crypto_secretstream_xchacha20poly1305_state()
|
||||
stream_header = sodium.crypto_secretstream_xchacha20poly1305_init_push(state, key)
|
||||
|
||||
chunks = []
|
||||
offset = 0
|
||||
total = len(plaintext)
|
||||
while offset < total:
|
||||
chunk = plaintext[offset : offset + CHUNK_SIZE]
|
||||
offset += len(chunk)
|
||||
is_last = offset >= total and not drop_final_tag
|
||||
tag = (
|
||||
sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL
|
||||
if is_last
|
||||
else sodium.crypto_secretstream_xchacha20poly1305_TAG_MESSAGE
|
||||
)
|
||||
chunks.append(sodium.crypto_secretstream_xchacha20poly1305_push(state, chunk, None, tag))
|
||||
ciphertext = b''.join(chunks)
|
||||
|
||||
digest = digest_override if digest_override is not None else hashlib.sha256(ciphertext).digest()
|
||||
header = HEADER_STRUCT.pack(
|
||||
magic,
|
||||
format_version,
|
||||
aead_id,
|
||||
kdf_id,
|
||||
salt,
|
||||
declared_memlimit if declared_memlimit is not None else memlimit,
|
||||
declared_opslimit if declared_opslimit is not None else opslimit,
|
||||
parallelism,
|
||||
digest,
|
||||
reserved_signature if reserved_signature is not None else bytes(RESERVED_SIGNATURE_SIZE),
|
||||
)
|
||||
return header + stream_header + ciphertext
|
||||
|
||||
|
||||
def make_bundle(
|
||||
zip_spec: ZipSpec | None = None,
|
||||
password: str = DEFAULT_PASSWORD,
|
||||
**header_kwargs: Any,
|
||||
) -> bytes:
|
||||
"""The common case: a well-formed tree, encrypted with the test cost parameters."""
|
||||
return encrypt_archive(build_zip(zip_spec), password, **header_kwargs)
|
||||
|
||||
|
||||
def tamper_ciphertext(bundle: bytes, offset_from_stream: int = 8) -> bytes:
|
||||
"""Flip one byte inside the chunk stream **and** repair the header digest.
|
||||
|
||||
Repairing the digest is what makes this a gate-3 test rather than a gate-2 one: the file is
|
||||
self-consistent and only the AEAD can tell it was altered.
|
||||
"""
|
||||
position = HEADER_SIZE + 24 + offset_from_stream
|
||||
altered = bytearray(bundle)
|
||||
altered[position] ^= 0xFF
|
||||
repaired = hashlib.sha256(bytes(altered[HEADER_SIZE + 24 :])).digest()
|
||||
altered[44:76] = repaired
|
||||
return bytes(altered)
|
||||
|
||||
|
||||
def truncate_bundle(bundle: bytes, drop_bytes: int = 32) -> bytes:
|
||||
"""Cut the tail off a bundle, leaving the header's digest describing bytes that are gone."""
|
||||
return bundle[:-drop_bytes]
|
||||
|
||||
|
||||
def build_envelope(password: str = DEFAULT_PASSWORD, key: bytes = DEFAULT_ENVELOPE_KEY) -> str:
|
||||
"""`base64(nonce || ciphertext || tag)` under AES-256-GCM, as the frontend will send it."""
|
||||
nonce = os.urandom(12)
|
||||
sealed = AESGCM(key).encrypt(nonce, password.encode('utf-8'), None)
|
||||
return base64.b64encode(nonce + sealed).decode('ascii')
|
||||
|
||||
|
||||
def envelope_key_b64(key: bytes = DEFAULT_ENVELOPE_KEY) -> str:
|
||||
"""The key as `IMPORT_PASSWORD_KEY` carries it."""
|
||||
return base64.b64encode(key).decode('ascii')
|
||||
Reference in New Issue
Block a user