From bc27757a739e5c0fb77697055ddf9373ab265eda Mon Sep 17 00:00:00 2001 From: Nikita Aksenov Date: Wed, 22 Jul 2026 19:45:43 +0300 Subject: [PATCH] feat: get camera snapshots from S3: untracked files --- src/services/snapshot_storage.py | 382 +++++++++++++++++++++++++++++++ tests/test_snapshot_storage.py | 248 ++++++++++++++++++++ 2 files changed, 630 insertions(+) create mode 100644 src/services/snapshot_storage.py create mode 100644 tests/test_snapshot_storage.py diff --git a/src/services/snapshot_storage.py b/src/services/snapshot_storage.py new file mode 100644 index 0000000..210e6a3 --- /dev/null +++ b/src/services/snapshot_storage.py @@ -0,0 +1,382 @@ +"""Read and decrypt detection artifacts stored as PTSNAP01 objects in S3.""" + +from __future__ import annotations + +import base64 +import binascii +import json +import os +import struct +from dataclasses import dataclass, field +from functools import lru_cache +from typing import Any, Mapping + +import boto3 +from botocore.exceptions import BotoCoreError, ClientError +from cryptography.exceptions import InvalidTag +from cryptography.hazmat.primitives.ciphers.aead import AESGCM + + +MAGIC = b"PTSNAP01" +FORMAT_VERSION = 1 +HEADER_LENGTH_STRUCT = struct.Struct(">I") +NONCE_SIZE = 12 +TAG_SIZE = 16 +DEFAULT_MAX_OBJECT_SIZE_BYTES = 25 * 1024 * 1024 +MAX_HEADER_SIZE_BYTES = 64 * 1024 +SUPPORTED_VARIANTS = {"raw", "annotated", "labels"} + + +class SnapshotStorageError(Exception): + """Base class for safe, user-facing storage error mapping.""" + + +class SnapshotNotConfiguredError(SnapshotStorageError): + pass + + +class SnapshotNotFoundError(SnapshotStorageError): + pass + + +class SnapshotUnavailableError(SnapshotStorageError): + pass + + +class SnapshotInvalidError(SnapshotStorageError): + pass + + +def _decode_key(encoded_key: str, variable_name: str) -> bytes: + try: + key = base64.b64decode(encoded_key, validate=True) + except (ValueError, binascii.Error) as exception: + raise ValueError(f"{variable_name} must be valid Base64") from exception + if len(key) != 32: + raise ValueError(f"{variable_name} must decode to exactly 32 bytes") + return key + + +@dataclass(frozen=True) +class SnapshotStorageSettings: + bucket: str + encryption_keys: Mapping[str, bytes] = field(repr=False) + prefix: str = "camera-snapshots" + endpoint_url: str | None = None + region_name: str | None = None + max_object_size_bytes: int = DEFAULT_MAX_OBJECT_SIZE_BYTES + + @classmethod + def from_environment( + cls, + environment: Mapping[str, str] | None = None, + ) -> "SnapshotStorageSettings": + env = environment if environment is not None else os.environ + bucket = env.get("SNAPSHOT_S3_BUCKET", "").strip() + if not bucket: + raise ValueError("SNAPSHOT_S3_BUCKET is required") + + keys: dict[str, bytes] = {} + current_key_id = env.get("SNAPSHOT_ENCRYPTION_KEY_ID", "").strip() + current_key_base64 = env.get("SNAPSHOT_ENCRYPTION_KEY_BASE64", "").strip() + if current_key_id and current_key_base64: + keys[current_key_id] = _decode_key( + current_key_base64, + "SNAPSHOT_ENCRYPTION_KEY_BASE64", + ) + elif current_key_id or current_key_base64: + raise ValueError( + "SNAPSHOT_ENCRYPTION_KEY_ID and SNAPSHOT_ENCRYPTION_KEY_BASE64 " + "must be configured together" + ) + + keyring_json = env.get("SNAPSHOT_DECRYPTION_KEYS_JSON", "").strip() + if keyring_json: + try: + keyring = json.loads(keyring_json) + except json.JSONDecodeError as exception: + raise ValueError("SNAPSHOT_DECRYPTION_KEYS_JSON must be valid JSON") from exception + if not isinstance(keyring, dict): + raise ValueError("SNAPSHOT_DECRYPTION_KEYS_JSON must be a JSON object") + for key_id, encoded_key in keyring.items(): + if not isinstance(key_id, str) or not key_id or not isinstance(encoded_key, str): + raise ValueError( + "SNAPSHOT_DECRYPTION_KEYS_JSON must map non-empty key IDs to Base64 strings" + ) + keys[key_id] = _decode_key(encoded_key, "SNAPSHOT_DECRYPTION_KEYS_JSON value") + + if not keys: + raise ValueError( + "At least one snapshot encryption key must be configured" + ) + + prefix = env.get("SNAPSHOT_S3_PREFIX", "camera-snapshots").strip("/") + if not prefix: + raise ValueError("SNAPSHOT_S3_PREFIX must not be empty") + + raw_max_size = env.get("SNAPSHOT_MAX_OBJECT_SIZE_BYTES", "").strip() + try: + max_size = int(raw_max_size) if raw_max_size else DEFAULT_MAX_OBJECT_SIZE_BYTES + except ValueError as exception: + raise ValueError("SNAPSHOT_MAX_OBJECT_SIZE_BYTES must be an integer") from exception + if max_size <= 0: + raise ValueError("SNAPSHOT_MAX_OBJECT_SIZE_BYTES must be positive") + + return cls( + bucket=bucket, + encryption_keys=keys, + prefix=prefix, + endpoint_url=env.get("SNAPSHOT_S3_ENDPOINT_URL") or None, + region_name=env.get("SNAPSHOT_S3_REGION") or None, + max_object_size_bytes=max_size, + ) + + +@dataclass(frozen=True) +class SnapshotReference: + bucket: str + object_key: str + variant: str + encryption_key_id: str | None = None + + +@dataclass(frozen=True) +class DecryptedSnapshot: + content: bytes = field(repr=False) + content_type: str + filename: str + captured_at: str + + +def snapshot_reference_from_metadata( + metadata: Any, + variant: str, +) -> SnapshotReference | None: + if variant not in SUPPORTED_VARIANTS or not isinstance(metadata, dict): + return None + snapshots = metadata.get("snapshots") + if not isinstance(snapshots, dict): + return None + value = snapshots.get(variant) + if not isinstance(value, dict): + return None + + bucket = value.get("bucket") + object_key = value.get("object_key") or value.get("key") + stored_variant = value.get("variant", variant) + key_id = value.get("encryption_key_id") or value.get("key_id") + if not isinstance(bucket, str) or not bucket: + return None + if not isinstance(object_key, str) or not object_key: + return None + if stored_variant != variant: + return None + if key_id is not None and (not isinstance(key_id, str) or not key_id): + return None + + return SnapshotReference( + bucket=bucket, + object_key=object_key, + variant=variant, + encryption_key_id=key_id, + ) + + +class S3SnapshotReader: + def __init__( + self, + settings: SnapshotStorageSettings, + *, + s3_client: Any | None = None, + ) -> None: + self.settings = settings + self.s3_client = s3_client or boto3.client( + "s3", + endpoint_url=settings.endpoint_url, + region_name=settings.region_name, + ) + + def read( + self, + reference: SnapshotReference, + *, + expected_camera_id: int, + expected_variant: str, + ) -> DecryptedSnapshot: + self._validate_reference(reference, expected_camera_id, expected_variant) + container = self._download(reference) + header_prefix, header, nonce, ciphertext = self._parse_container(container) + + key_id = reference.encryption_key_id or header.get("key_id") + if not isinstance(key_id, str) or key_id not in self.settings.encryption_keys: + raise SnapshotNotConfiguredError("Snapshot encryption key is unavailable") + + try: + plaintext = AESGCM(self.settings.encryption_keys[key_id]).decrypt( + nonce, + ciphertext, + header_prefix, + ) + except InvalidTag as exception: + raise SnapshotInvalidError("Snapshot integrity check failed") from exception + + content_type, captured_at = self._validate_header( + header, + key_id=key_id, + expected_camera_id=expected_camera_id, + expected_variant=expected_variant, + ) + if expected_variant == "labels": + try: + plaintext.decode("utf-8") + except UnicodeDecodeError as exception: + raise SnapshotInvalidError("Snapshot labels are not valid UTF-8") from exception + + filename = "labels.txt" if expected_variant == "labels" else f"{expected_variant}.jpg" + return DecryptedSnapshot( + content=plaintext, + content_type=content_type, + filename=filename, + captured_at=captured_at, + ) + + def _validate_reference( + self, + reference: SnapshotReference, + expected_camera_id: int, + expected_variant: str, + ) -> None: + if expected_variant not in SUPPORTED_VARIANTS or reference.variant != expected_variant: + raise SnapshotInvalidError("Snapshot variant does not match the request") + if reference.bucket != self.settings.bucket: + raise SnapshotInvalidError("Snapshot bucket does not match API configuration") + expected_prefix = f"{self.settings.prefix}/camera-{expected_camera_id}/" + if not reference.object_key.startswith(expected_prefix): + raise SnapshotInvalidError("Snapshot object key does not match the detection camera") + + def _download(self, reference: SnapshotReference) -> bytes: + try: + response = self.s3_client.get_object( + Bucket=reference.bucket, + Key=reference.object_key, + ) + content_length = response.get("ContentLength") + if ( + isinstance(content_length, int) + and content_length > self.settings.max_object_size_bytes + ): + raise SnapshotInvalidError("Encrypted snapshot is too large") + body = response["Body"] + try: + container = body.read(self.settings.max_object_size_bytes + 1) + finally: + close = getattr(body, "close", None) + if callable(close): + close() + except ClientError as exception: + error_code = str(exception.response.get("Error", {}).get("Code", "")) + if error_code in {"404", "NoSuchKey", "NotFound"}: + raise SnapshotNotFoundError("Snapshot object not found") from exception + raise SnapshotUnavailableError("S3 snapshot storage is unavailable") from exception + except (BotoCoreError, KeyError, OSError) as exception: + raise SnapshotUnavailableError("S3 snapshot storage is unavailable") from exception + + if not isinstance(container, bytes): + container = bytes(container) + if len(container) > self.settings.max_object_size_bytes: + raise SnapshotInvalidError("Encrypted snapshot is too large") + return container + + @staticmethod + def _parse_container( + container: bytes, + ) -> tuple[bytes, dict[str, Any], bytes, bytes]: + minimum_size = len(MAGIC) + HEADER_LENGTH_STRUCT.size + NONCE_SIZE + TAG_SIZE + if len(container) < minimum_size or not container.startswith(MAGIC): + raise SnapshotInvalidError("Invalid encrypted snapshot format") + + header_length_offset = len(MAGIC) + header_offset = header_length_offset + HEADER_LENGTH_STRUCT.size + (header_length,) = HEADER_LENGTH_STRUCT.unpack_from(container, header_length_offset) + if header_length <= 0 or header_length > MAX_HEADER_SIZE_BYTES: + raise SnapshotInvalidError("Invalid encrypted snapshot header length") + nonce_offset = header_offset + header_length + ciphertext_offset = nonce_offset + NONCE_SIZE + if ciphertext_offset + TAG_SIZE > len(container): + raise SnapshotInvalidError("Encrypted snapshot is truncated") + + try: + header = json.loads(container[header_offset:nonce_offset].decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exception: + raise SnapshotInvalidError("Invalid encrypted snapshot header") from exception + if not isinstance(header, dict): + raise SnapshotInvalidError("Invalid encrypted snapshot header") + + return ( + container[:nonce_offset], + header, + container[nonce_offset:ciphertext_offset], + container[ciphertext_offset:], + ) + + @staticmethod + def _validate_header( + header: Mapping[str, Any], + *, + key_id: str, + expected_camera_id: int, + expected_variant: str, + ) -> tuple[str, str]: + if header.get("format_version") != FORMAT_VERSION: + raise SnapshotInvalidError("Unsupported snapshot format version") + if header.get("algorithm") != "AES-256-GCM": + raise SnapshotInvalidError("Unsupported snapshot encryption algorithm") + if header.get("key_id") != key_id: + raise SnapshotInvalidError("Snapshot key ID does not match metadata") + if header.get("camera_id") != expected_camera_id: + raise SnapshotInvalidError("Snapshot camera ID does not match the detection") + if header.get("variant") != expected_variant: + raise SnapshotInvalidError("Snapshot variant does not match the request") + + captured_at = header.get("captured_at") + content_type = header.get("content_type") + if not isinstance(captured_at, str) or not captured_at: + raise SnapshotInvalidError("Snapshot capture time is missing") + if not isinstance(content_type, str): + raise SnapshotInvalidError("Snapshot content type is missing") + + if expected_variant in {"raw", "annotated"}: + if content_type != "image/jpeg": + raise SnapshotInvalidError("Snapshot image content type is invalid") + else: + if not content_type.startswith("text/plain"): + raise SnapshotInvalidError("Snapshot labels content type is invalid") + if header.get("annotation_format") != "yolo-v12-detection": + raise SnapshotInvalidError("Snapshot labels format is invalid") + + return content_type, captured_at + + +@lru_cache(maxsize=1) +def get_snapshot_reader() -> S3SnapshotReader: + try: + settings = SnapshotStorageSettings.from_environment() + except ValueError as exception: + raise SnapshotNotConfiguredError(str(exception)) from exception + return S3SnapshotReader(settings) + + +def read_snapshot_from_metadata( + metadata: Any, + *, + camera_id: int, + variant: str, +) -> DecryptedSnapshot: + reference = snapshot_reference_from_metadata(metadata, variant) + if reference is None: + raise SnapshotNotFoundError("Detection artifact metadata is missing") + return get_snapshot_reader().read( + reference, + expected_camera_id=camera_id, + expected_variant=variant, + ) diff --git a/tests/test_snapshot_storage.py b/tests/test_snapshot_storage.py new file mode 100644 index 0000000..41d0038 --- /dev/null +++ b/tests/test_snapshot_storage.py @@ -0,0 +1,248 @@ +from __future__ import annotations + +import base64 +import io +import json +import struct +import unittest + +from botocore.exceptions import ClientError +from cryptography.hazmat.primitives.ciphers.aead import AESGCM + +from src.services.snapshot_storage import ( + MAGIC, + MAX_HEADER_SIZE_BYTES, + S3SnapshotReader, + SnapshotInvalidError, + SnapshotReference, + SnapshotStorageSettings, + SnapshotUnavailableError, + snapshot_reference_from_metadata, +) + + +TEST_KEY = bytes(range(32)) +TEST_KEY_ID = "snapshot-key-test" +TEST_BUCKET = "snapshots" +TEST_PREFIX = "camera-snapshots" +TEST_OBJECT_KEY = ( + "camera-snapshots/camera-17/2026/07/22/" + "20260722T080910.123456Z_test/raw.jpg.aesgcm" +) + + +def build_container( + plaintext: bytes, + *, + camera_id: int = 17, + variant: str = "raw", + content_type: str = "image/jpeg", +) -> bytes: + header = { + "algorithm": "AES-256-GCM", + "camera_id": camera_id, + "captured_at": "2026-07-22T08:09:10.123456Z", + "content_type": content_type, + "format_version": 1, + "key_id": TEST_KEY_ID, + "variant": variant, + } + if variant == "labels": + header["annotation_format"] = "yolo-v12-detection" + header_bytes = json.dumps( + header, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + prefix = MAGIC + struct.pack(">I", len(header_bytes)) + header_bytes + nonce = bytes(range(12)) + return prefix + nonce + AESGCM(TEST_KEY).encrypt(nonce, plaintext, prefix) + + +class TrackingBody(io.BytesIO): + def __init__(self, content: bytes): + super().__init__(content) + self.was_closed = False + + def close(self): + self.was_closed = True + super().close() + + +class FakeS3Client: + def __init__(self, content: bytes): + self.content = content + self.calls: list[dict[str, str]] = [] + self.body: TrackingBody | None = None + + def get_object(self, **kwargs): + self.calls.append(kwargs) + self.body = TrackingBody(self.content) + return {"Body": self.body, "ContentLength": len(self.content)} + + +class ErrorS3Client: + def __init__(self, code: str): + self.code = code + + def get_object(self, **kwargs): + raise ClientError( + {"Error": {"Code": self.code, "Message": self.code}}, + "GetObject", + ) + + +class SnapshotStorageSettingsTests(unittest.TestCase): + def test_loads_current_and_rotated_keys(self): + old_key = bytes(reversed(range(32))) + settings = SnapshotStorageSettings.from_environment({ + "SNAPSHOT_S3_BUCKET": TEST_BUCKET, + "SNAPSHOT_ENCRYPTION_KEY_ID": TEST_KEY_ID, + "SNAPSHOT_ENCRYPTION_KEY_BASE64": base64.b64encode(TEST_KEY).decode(), + "SNAPSHOT_DECRYPTION_KEYS_JSON": json.dumps({ + "snapshot-key-old": base64.b64encode(old_key).decode(), + }), + }) + + self.assertEqual(TEST_KEY, settings.encryption_keys[TEST_KEY_ID]) + self.assertEqual(old_key, settings.encryption_keys["snapshot-key-old"]) + + def test_extracts_nested_snapshot_reference(self): + metadata = { + "snapshots": { + "raw": { + "bucket": TEST_BUCKET, + "object_key": TEST_OBJECT_KEY, + "variant": "raw", + "encryption_key_id": TEST_KEY_ID, + } + } + } + + reference = snapshot_reference_from_metadata(metadata, "raw") + + self.assertEqual( + SnapshotReference(TEST_BUCKET, TEST_OBJECT_KEY, "raw", TEST_KEY_ID), + reference, + ) + self.assertIsNone(snapshot_reference_from_metadata(metadata, "annotated")) + + +class S3SnapshotReaderTests(unittest.TestCase): + def make_reader(self, container: bytes) -> tuple[S3SnapshotReader, FakeS3Client]: + client = FakeS3Client(container) + settings = SnapshotStorageSettings( + bucket=TEST_BUCKET, + encryption_keys={TEST_KEY_ID: TEST_KEY}, + prefix=TEST_PREFIX, + ) + return S3SnapshotReader(settings, s3_client=client), client + + def test_downloads_authenticates_and_decrypts_image(self): + plaintext = b"jpeg bytes" + reader, client = self.make_reader(build_container(plaintext)) + reference = SnapshotReference( + TEST_BUCKET, + TEST_OBJECT_KEY, + "raw", + TEST_KEY_ID, + ) + + result = reader.read(reference, expected_camera_id=17, expected_variant="raw") + + self.assertEqual(plaintext, result.content) + self.assertEqual("image/jpeg", result.content_type) + self.assertEqual("raw.jpg", result.filename) + self.assertEqual( + [{"Bucket": TEST_BUCKET, "Key": TEST_OBJECT_KEY}], + client.calls, + ) + self.assertIsNotNone(client.body) + self.assertTrue(client.body.was_closed) + + def test_downloads_yolov12_labels(self): + object_key = TEST_OBJECT_KEY.replace("raw.jpg", "labels.txt") + plaintext = b"0 0.500000 0.500000 0.250000 0.500000\n" + reader, _ = self.make_reader(build_container( + plaintext, + variant="labels", + content_type="text/plain; charset=utf-8", + )) + reference = SnapshotReference(TEST_BUCKET, object_key, "labels", TEST_KEY_ID) + + result = reader.read(reference, expected_camera_id=17, expected_variant="labels") + + self.assertEqual(plaintext, result.content) + self.assertEqual("text/plain; charset=utf-8", result.content_type) + self.assertEqual("labels.txt", result.filename) + + def test_rejects_tampered_ciphertext(self): + container = bytearray(build_container(b"jpeg bytes")) + container[-1] ^= 1 + reader, _ = self.make_reader(bytes(container)) + reference = SnapshotReference( + TEST_BUCKET, + TEST_OBJECT_KEY, + "raw", + TEST_KEY_ID, + ) + + with self.assertRaises(SnapshotInvalidError): + reader.read(reference, expected_camera_id=17, expected_variant="raw") + + def test_rejects_oversized_header_before_json_parsing(self): + container = ( + MAGIC + + struct.pack(">I", MAX_HEADER_SIZE_BYTES + 1) + + bytes(12 + 16) + ) + + with self.assertRaises(SnapshotInvalidError): + S3SnapshotReader._parse_container(container) + + def test_rejects_non_utf8_labels(self): + object_key = TEST_OBJECT_KEY.replace("raw.jpg", "labels.txt") + reader, _ = self.make_reader(build_container( + b"\xff\xfe", + variant="labels", + content_type="text/plain; charset=utf-8", + )) + reference = SnapshotReference(TEST_BUCKET, object_key, "labels", TEST_KEY_ID) + + with self.assertRaises(SnapshotInvalidError): + reader.read(reference, expected_camera_id=17, expected_variant="labels") + + def test_maps_missing_bucket_to_storage_unavailable(self): + settings = SnapshotStorageSettings( + bucket=TEST_BUCKET, + encryption_keys={TEST_KEY_ID: TEST_KEY}, + prefix=TEST_PREFIX, + ) + reader = S3SnapshotReader(settings, s3_client=ErrorS3Client("NoSuchBucket")) + reference = SnapshotReference( + TEST_BUCKET, + TEST_OBJECT_KEY, + "raw", + TEST_KEY_ID, + ) + + with self.assertRaises(SnapshotUnavailableError): + reader.read(reference, expected_camera_id=17, expected_variant="raw") + + def test_rejects_object_key_from_another_camera_before_s3_request(self): + reader, client = self.make_reader(build_container(b"jpeg bytes")) + reference = SnapshotReference( + TEST_BUCKET, + TEST_OBJECT_KEY.replace("camera-17", "camera-99"), + "raw", + TEST_KEY_ID, + ) + + with self.assertRaises(SnapshotInvalidError): + reader.read(reference, expected_camera_id=17, expected_variant="raw") + self.assertEqual([], client.calls) + + +if __name__ == "__main__": + unittest.main()