vitor-aignosi
d257e70465
QTZPOC-20: Refactor test inputs into named variables for improved readability and debugging.
486 lines
16 KiB
Python
486 lines
16 KiB
Python
"""The seven gates: what each one rejects, in what order, and what it leaves on disk."""
|
|
|
|
import json
|
|
import random
|
|
from pathlib import Path
|
|
|
|
import nacl.bindings as sodium
|
|
import pytest
|
|
|
|
from laborious.utils.bundle import format as fmt
|
|
from laborious.utils.bundle.reader import BundleGateError, BundleLimits, BundleReader
|
|
from laborious.utils.bundle.steps import ImportStep
|
|
from tests.helpers import bundle_factory as factory
|
|
|
|
OBJECT_KEY = 'imported_models/factory_model_v2.sientia'
|
|
|
|
|
|
@pytest.fixture
|
|
def reader() -> BundleReader:
|
|
return BundleReader()
|
|
|
|
|
|
def _bundle_file(work_dir: Path, payload: bytes) -> Path:
|
|
target = work_dir / 'download.sientia'
|
|
target.write_bytes(payload)
|
|
return target
|
|
|
|
|
|
def _open(reader: BundleReader, work_dir: Path, payload: bytes, password: str | None = None):
|
|
return reader.open(
|
|
_bundle_file(work_dir, payload),
|
|
password if password is not None else factory.DEFAULT_PASSWORD,
|
|
work_dir / 'extract',
|
|
work_dir,
|
|
)
|
|
|
|
|
|
def _extracted_files(root: Path) -> list[str]:
|
|
if not root.exists():
|
|
return []
|
|
return sorted(str(p.relative_to(root)) for p in root.rglob('*') if p.is_file())
|
|
|
|
|
|
# ------------------------------------------------------------------- gate 1
|
|
|
|
|
|
def test_gate_1_accepts_a_well_shaped_object(reader):
|
|
reader.check_object_limits(size=1024, key=OBJECT_KEY)
|
|
|
|
|
|
def test_gate_1_rejects_an_oversized_object(reader):
|
|
limited = BundleReader(BundleLimits(max_object_bytes=1000))
|
|
|
|
with pytest.raises(BundleGateError) as error:
|
|
limited.check_object_limits(size=1001, key=OBJECT_KEY)
|
|
|
|
assert error.value.gate == 1
|
|
assert error.value.step is ImportStep.DOWNLOAD
|
|
|
|
|
|
def test_gate_1_rejects_an_empty_object(reader):
|
|
with pytest.raises(BundleGateError, match='empty'):
|
|
reader.check_object_limits(size=0, key=OBJECT_KEY)
|
|
|
|
|
|
def test_gate_1_rejects_a_key_outside_the_import_prefix(reader):
|
|
with pytest.raises(BundleGateError, match='import prefix'):
|
|
reader.check_object_limits(size=10, key='other_place/model.sientia')
|
|
|
|
|
|
def test_gate_1_rejects_a_key_with_the_wrong_suffix(reader):
|
|
with pytest.raises(BundleGateError, match='.sientia'):
|
|
reader.check_object_limits(size=10, key='imported_models/model.zip')
|
|
|
|
|
|
# ------------------------------------------------------------- golden bundle
|
|
|
|
|
|
def test_opens_the_producer_fixture(golden_bundle, golden_password, work_dir, reader):
|
|
"""The producer's own bytes, at the producer's own KDF cost."""
|
|
opened = reader.open(golden_bundle, golden_password, work_dir / 'extract', work_dir)
|
|
|
|
assert opened.metadata['model_name'] == 'fixture_model'
|
|
assert opened.metadata['experiment_name'] == 'fixture_experiment'
|
|
assert set(opened.metadata) == {
|
|
'model_name',
|
|
'experiment_name',
|
|
'run_id',
|
|
'model_version',
|
|
'model_project',
|
|
'export_timestamp',
|
|
}
|
|
assert opened.parameters['target_variable'] == 'SE'
|
|
assert opened.header is not None
|
|
assert opened.digest == opened.header.ciphertext_digest.hex()
|
|
assert _extracted_files(opened.extracted_dir) == [
|
|
'artifacts/data_model/MLmodel',
|
|
'artifacts/data_model/conda.yaml',
|
|
'artifacts/data_model/model.pkl',
|
|
'artifacts/data_model/python_env.yaml',
|
|
'artifacts/data_model/requirements.txt',
|
|
'artifacts/model_card/model_card.json',
|
|
'artifacts/prediction_model/MLmodel',
|
|
'artifacts/prediction_model/conda.yaml',
|
|
'artifacts/prediction_model/model.pkl',
|
|
'artifacts/prediction_model/python_env.yaml',
|
|
'artifacts/prediction_model/requirements.txt',
|
|
'artifacts/raw_data.csv',
|
|
'artifacts/test_data.csv',
|
|
'artifacts/train_data.csv',
|
|
'metadata.json',
|
|
]
|
|
assert list(work_dir.glob('*.zip')) == []
|
|
|
|
|
|
def test_fixture_metadata_matches_its_committed_plaintext_copy(
|
|
golden_bundle, golden_password, work_dir, reader
|
|
):
|
|
from tests.laborious.utils.bundle.conftest import GOLDEN_METADATA
|
|
|
|
opened = reader.open(golden_bundle, golden_password, work_dir / 'extract', work_dir)
|
|
committed = json.loads(GOLDEN_METADATA.read_text(encoding='utf-8'))
|
|
|
|
assert opened.metadata == committed['metadata']
|
|
assert opened.parameters == committed['parameters']
|
|
|
|
|
|
# ------------------------------------------------------------------- gate 2
|
|
|
|
|
|
def test_gate_2_rejects_a_wrong_magic_before_deriving(work_dir, reader, monkeypatch):
|
|
payload = factory.make_bundle(magic=b'NOTSIENT')
|
|
monkeypatch.setattr(
|
|
sodium, 'crypto_pwhash_alg', _never_called('key derivation must not run at gate 2')
|
|
)
|
|
|
|
with pytest.raises(BundleGateError) as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 2
|
|
assert error.value.step is ImportStep.DOWNLOAD
|
|
|
|
|
|
def test_gate_2_rejects_kdf_parameters_above_the_ceiling_without_allocating(
|
|
work_dir, reader, monkeypatch
|
|
):
|
|
payload = factory.make_bundle(declared_memlimit=512 * 1024 * 1024)
|
|
monkeypatch.setattr(
|
|
sodium, 'crypto_pwhash_alg', _never_called('the ceiling must be enforced before deriving')
|
|
)
|
|
|
|
with pytest.raises(BundleGateError) as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 2
|
|
|
|
|
|
def test_gate_2_rejects_a_digest_mismatch(work_dir, reader):
|
|
payload = factory.make_bundle(digest_override=bytes(32))
|
|
|
|
with pytest.raises(BundleGateError) as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 2
|
|
|
|
|
|
def test_gate_2_rejects_a_truncated_transfer(work_dir, reader):
|
|
payload = factory.truncate_bundle(factory.make_bundle())
|
|
|
|
with pytest.raises(BundleGateError) as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 2
|
|
|
|
|
|
def _never_called(message: str):
|
|
def _raise(*_args, **_kwargs):
|
|
raise AssertionError(message)
|
|
|
|
return _raise
|
|
|
|
|
|
# ------------------------------------------------------------------- gate 3
|
|
|
|
AMBIGUOUS_REASON = 'wrong password, or the file is damaged or was altered'
|
|
|
|
|
|
def test_gate_3_wrong_password(work_dir, reader):
|
|
payload = factory.make_bundle()
|
|
|
|
with pytest.raises(BundleGateError) as error:
|
|
_open(reader, work_dir, payload, password='Wrong1Password!Here')
|
|
|
|
assert error.value.gate == 3
|
|
assert error.value.step is ImportStep.DECRYPTION
|
|
assert AMBIGUOUS_REASON in str(error.value)
|
|
|
|
|
|
def test_gate_3_altered_ciphertext_reads_the_same_as_a_wrong_password(work_dir, reader):
|
|
tampered = factory.tamper_ciphertext(factory.make_bundle())
|
|
|
|
with pytest.raises(BundleGateError) as altered:
|
|
_open(reader, work_dir, tampered)
|
|
|
|
wrong_password_payload = factory.make_bundle()
|
|
with pytest.raises(BundleGateError) as wrong_password:
|
|
_open(reader, work_dir, wrong_password_payload, password='Wrong1Password!Here')
|
|
|
|
assert str(altered.value) == str(wrong_password.value)
|
|
assert altered.value.gate == wrong_password.value.gate == 3
|
|
|
|
|
|
def test_gate_3_stream_without_a_final_tag_is_rejected(work_dir, reader):
|
|
payload = factory.make_bundle(drop_final_tag=True)
|
|
|
|
with pytest.raises(BundleGateError) as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 3
|
|
assert AMBIGUOUS_REASON in str(error.value)
|
|
|
|
|
|
def test_gate_3_leaves_no_decrypted_zip_behind(work_dir, reader):
|
|
payload = factory.make_bundle()
|
|
|
|
with pytest.raises(BundleGateError):
|
|
_open(reader, work_dir, payload, password='Wrong1Password!Here')
|
|
|
|
assert list(work_dir.glob('*.zip')) == []
|
|
|
|
|
|
def test_gate_3_normalises_the_password_to_nfc(work_dir, reader):
|
|
composed = 'Café1Model!Test'
|
|
decomposed = 'Café1Model!Test'
|
|
payload = factory.make_bundle(password=composed)
|
|
|
|
opened = _open(reader, work_dir, payload, password=decomposed)
|
|
|
|
assert opened.metadata['model_name'] == 'factory_model'
|
|
|
|
|
|
def test_gate_3_does_not_trim_the_password(work_dir, reader):
|
|
payload = factory.make_bundle(password='Trailing1Space!Test ')
|
|
|
|
with pytest.raises(BundleGateError):
|
|
_open(reader, work_dir, payload, password='Trailing1Space!Test')
|
|
|
|
|
|
def test_multi_chunk_plaintext_round_trips(work_dir, reader):
|
|
"""More than 64 KiB of plaintext, so `MESSAGE` and `FINAL` framing are both exercised."""
|
|
spec = factory.ZipSpec()
|
|
# Incompressible on purpose: a compressible filler would deflate below one chunk and the test
|
|
# would silently stop covering multi-chunk framing.
|
|
spec.files['artifacts/prediction_model/model.pkl'] = random.Random( # noqa: S311 - filler
|
|
20260825
|
|
).randbytes(200 * 1024)
|
|
payload = factory.make_bundle(spec)
|
|
assert len(payload) > fmt.CHUNK_SIZE + fmt.ABYTES + fmt.DIGEST_OFFSET
|
|
|
|
opened = _open(reader, work_dir, payload)
|
|
|
|
extracted = opened.extracted_dir / 'artifacts' / 'prediction_model' / 'model.pkl'
|
|
assert extracted.read_bytes() == spec.files['artifacts/prediction_model/model.pkl']
|
|
|
|
|
|
# ------------------------------------------------------------------- gate 4
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
('spec_kwargs', 'match'),
|
|
[
|
|
({'bomb_bytes': 40 * 1024 * 1024}, 'compression ratio'),
|
|
({'traversal_entry': 'artifacts/../../etc/passwd'}, 'escaping'),
|
|
({'absolute_entry': '/etc/passwd'}, 'absolute'),
|
|
({'symlinks': {'artifacts/link': '/etc/passwd'}}, 'symbolic link'),
|
|
],
|
|
)
|
|
def test_gate_4_rejects_hostile_central_directories(work_dir, reader, spec_kwargs, match):
|
|
payload = factory.make_bundle(factory.ZipSpec(**spec_kwargs))
|
|
|
|
with pytest.raises(BundleGateError, match=match) as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 4
|
|
assert error.value.step is ImportStep.ARCHIVE_INSPECTION
|
|
assert _extracted_files(work_dir / 'extract') == []
|
|
assert not (work_dir / 'extract').exists()
|
|
|
|
|
|
def test_gate_4_rejects_too_many_entries(work_dir):
|
|
limited = BundleReader(BundleLimits(max_entries=3))
|
|
payload = factory.make_bundle()
|
|
|
|
with pytest.raises(BundleGateError, match='entries') as error:
|
|
_open(limited, work_dir, payload)
|
|
|
|
assert error.value.gate == 4
|
|
|
|
|
|
def test_gate_4_rejects_an_oversized_uncompressed_total(work_dir):
|
|
limited = BundleReader(BundleLimits(max_uncompressed_bytes=64))
|
|
payload = factory.make_bundle()
|
|
|
|
with pytest.raises(BundleGateError, match='uncompressed') as error:
|
|
_open(limited, work_dir, payload)
|
|
|
|
assert error.value.gate == 4
|
|
|
|
|
|
def test_gate_4_rejects_content_that_is_not_an_archive(work_dir, reader):
|
|
payload = factory.encrypt_archive(b'not a zip at all')
|
|
|
|
with pytest.raises(BundleGateError, match='readable archive') as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 4
|
|
|
|
|
|
# ------------------------------------------------------------------- gate 5
|
|
|
|
|
|
def test_gate_5_refuses_a_pre_existing_extraction_root(work_dir, reader):
|
|
(work_dir / 'extract').mkdir()
|
|
payload = factory.make_bundle()
|
|
|
|
with pytest.raises(BundleGateError, match='already exists') as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 5
|
|
assert error.value.step is ImportStep.EXTRACTION
|
|
|
|
|
|
def test_gate_5_extraction_writes_plain_files(work_dir, reader):
|
|
opened = _open(reader, work_dir, factory.make_bundle())
|
|
|
|
for path in opened.extracted_dir.rglob('*'):
|
|
if path.is_file():
|
|
assert path.stat().st_mode & 0o111 == 0, path
|
|
|
|
|
|
# ------------------------------------------------------------------- gate 6
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
('spec', 'match'),
|
|
[
|
|
(factory.ZipSpec(include_metadata=False), 'metadata document'),
|
|
(
|
|
factory.ZipSpec(
|
|
document=factory.build_metadata_document(drop_top_level_key='parameters')
|
|
),
|
|
'parameters and metadata blocks',
|
|
),
|
|
(
|
|
factory.ZipSpec(
|
|
document=factory.build_metadata_document(drop_metadata_field='experiment_name')
|
|
),
|
|
'experiment_name',
|
|
),
|
|
(
|
|
factory.ZipSpec(
|
|
document=factory.build_metadata_document(drop_parameter='target_variable')
|
|
),
|
|
'target_variable',
|
|
),
|
|
(
|
|
factory.ZipSpec(
|
|
document=factory.build_metadata_document(metadata={'experiment_name': 'x' * 51})
|
|
),
|
|
'experiment name',
|
|
),
|
|
(
|
|
factory.ZipSpec(
|
|
document=factory.build_metadata_document(metadata={'experiment_name': 'ab'})
|
|
),
|
|
'experiment name',
|
|
),
|
|
],
|
|
)
|
|
def test_gate_6_rejects_a_malformed_metadata_document(work_dir, reader, spec, match):
|
|
payload = factory.make_bundle(spec)
|
|
|
|
with pytest.raises(BundleGateError, match=match) as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 6
|
|
assert error.value.step is ImportStep.STRUCTURE_VALIDATION
|
|
assert not (work_dir / 'extract').exists()
|
|
|
|
|
|
def test_gate_6_rejects_unreadable_metadata_json(work_dir, reader):
|
|
spec = factory.ZipSpec(include_metadata=False)
|
|
spec.extra_files['metadata.json'] = b'{not json'
|
|
payload = factory.make_bundle(spec)
|
|
|
|
with pytest.raises(BundleGateError, match='readable JSON') as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 6
|
|
|
|
|
|
def test_gate_6_rejects_a_missing_artifacts_directory(work_dir, reader):
|
|
spec = factory.ZipSpec(files={})
|
|
payload = factory.make_bundle(spec)
|
|
|
|
with pytest.raises(BundleGateError, match='artifacts directory') as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 6
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
('dropped', 'match'),
|
|
[
|
|
('artifacts/data_model/MLmodel', 'data_model artifacts are incomplete'),
|
|
('artifacts/prediction_model/model.pkl', 'prediction_model artifacts are incomplete'),
|
|
],
|
|
)
|
|
def test_gate_6_rejects_an_incomplete_model_directory(work_dir, reader, dropped, match):
|
|
spec = factory.ZipSpec()
|
|
spec.files.pop(dropped)
|
|
payload = factory.make_bundle(spec)
|
|
|
|
with pytest.raises(BundleGateError, match=match) as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 6
|
|
|
|
|
|
def test_gate_6_rejects_a_bundle_without_data_model(work_dir, reader):
|
|
spec = factory.ZipSpec()
|
|
for name in list(spec.files):
|
|
if name.startswith('artifacts/data_model/'):
|
|
spec.files.pop(name)
|
|
payload = factory.make_bundle(spec)
|
|
|
|
with pytest.raises(BundleGateError, match='no data_model artifacts') as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 6
|
|
|
|
|
|
# ------------------------------------------------------------------- gate 7
|
|
|
|
|
|
@pytest.mark.parametrize('name', ['artifacts/run.sh', 'artifacts/tool.bin', 'artifacts/mod.so'])
|
|
def test_gate_7_rejects_content_a_model_export_should_not_carry(work_dir, reader, name):
|
|
spec = factory.ZipSpec()
|
|
spec.extra_files[name] = b'#!/bin/sh\n'
|
|
payload = factory.make_bundle(spec)
|
|
|
|
with pytest.raises(BundleGateError, match='should not carry') as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 7
|
|
assert error.value.step is ImportStep.CONTENT_POLICY
|
|
assert not (work_dir / 'extract').exists()
|
|
|
|
|
|
# ------------------------------------------------------------------ ordering
|
|
|
|
|
|
def test_the_earlier_gate_wins(work_dir, reader):
|
|
"""A bundle that would fail gate 4 and gate 6 must report gate 4, having extracted nothing."""
|
|
spec = factory.ZipSpec(
|
|
traversal_entry='artifacts/../escape.csv',
|
|
document=factory.build_metadata_document(drop_parameter='target_variable'),
|
|
)
|
|
payload = factory.make_bundle(spec)
|
|
|
|
with pytest.raises(BundleGateError) as error:
|
|
_open(reader, work_dir, payload)
|
|
|
|
assert error.value.gate == 4
|
|
assert _extracted_files(work_dir / 'extract') == []
|
|
|
|
|
|
def test_header_is_checked_before_the_password(work_dir, reader):
|
|
"""A bundle with a bad header and a wrong password reports the header, not the password."""
|
|
payload = factory.make_bundle(magic=b'NOTSIENT')
|
|
|
|
with pytest.raises(BundleGateError) as error:
|
|
_open(reader, work_dir, payload, password='Nope1!Nope1!')
|
|
|
|
assert error.value.gate == 2
|