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.
273 lines
10 KiB
Python
273 lines
10 KiB
Python
"""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')
|