Add safe archive preview and extraction workflows

2026-07-31 02:48:56 +02:00
parent 159a012833
commit b752dea610
20 changed files with 1961 additions and 182 deletions
@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+import io
+import tarfile
+import unittest
+import zipfile
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pyzipper
+
+from govoplan_files.backend.storage.archives import (
+    ArchivePasswordError,
+    extract_archive_upload,
+    inspect_archive,
+)
+from govoplan_files.backend.storage.common import FileStorageError
+
+
+def _zip_bytes(entries: dict[str, bytes]) -> bytes:
+    target = io.BytesIO()
+    with zipfile.ZipFile(target, "w", compression=zipfile.ZIP_DEFLATED) as archive:
+        for path, data in entries.items():
+            archive.writestr(path, data)
+    return target.getvalue()
+
+
+def _tar_bytes(entries: dict[str, bytes], mode: str = "w:gz") -> bytes:
+    target = io.BytesIO()
+    with tarfile.open(fileobj=target, mode=mode) as archive:
+        for path, data in entries.items():
+            info = tarfile.TarInfo(path)
+            info.size = len(data)
+            archive.addfile(info, io.BytesIO(data))
+    return target.getvalue()
+
+
+def _encrypted_zip_bytes(password: str) -> bytes:
+    target = io.BytesIO()
+    with pyzipper.AESZipFile(
+        target,
+        "w",
+        compression=zipfile.ZIP_DEFLATED,
+        encryption=pyzipper.WZ_AES,
+    ) as archive:
+        archive.setpassword(password.encode("utf-8"))
+        archive.writestr("secure/report.txt", b"classified")
+    return target.getvalue()
+
+
+class ArchiveInspectionTests(unittest.TestCase):
+    def test_zip_preview_derives_folders_and_sizes(self) -> None:
+        inspection = inspect_archive(
+            _zip_bytes(
+                {
+                    "department/one.txt": b"one",
+                    "department/nested/two.csv": b"two",
+                }
+            ),
+            filename="monthly.zip",
+        )
+
+        self.assertEqual("zip", inspection.archive_format)
+        self.assertEqual(2, inspection.file_count)
+        self.assertEqual(2, inspection.directory_count)
+        self.assertEqual(6, inspection.expanded_size_bytes)
+        self.assertEqual(
+            [
+                "department",
+                "department/nested",
+                "department/nested/two.csv",
+                "department/one.txt",
+            ],
+            [entry.path for entry in inspection.entries],
+        )
+
+    def test_tar_compression_variants_are_inspected(self) -> None:
+        for filename, mode in (
+            ("monthly.tar", "w"),
+            ("monthly.tar.gz", "w:gz"),
+            ("monthly.tar.bz2", "w:bz2"),
+            ("monthly.tar.xz", "w:xz"),
+        ):
+            with self.subTest(filename=filename):
+                inspection = inspect_archive(
+                    _tar_bytes({"report.txt": b"report"}, mode=mode),
+                    filename=filename,
+                )
+                self.assertEqual(1, inspection.file_count)
+                self.assertEqual(filename.removeprefix("monthly."), inspection.archive_format)
+
+    def test_unsafe_member_path_is_rejected(self) -> None:
+        with self.assertRaisesRegex(FileStorageError, "Unsafe archive member"):
+            inspect_archive(
+                _zip_bytes({"../escape.txt": b"no"}),
+                filename="unsafe.zip",
+            )
+
+    def test_archive_bomb_ratio_is_rejected(self) -> None:
+        with self.assertRaisesRegex(FileStorageError, "expansion ratio"):
+            inspect_archive(
+                _zip_bytes({"zeros.bin": b"\0" * 50_000}),
+                filename="bomb.zip",
+                max_expansion_ratio=2,
+            )
+
+    def test_encrypted_zip_reports_and_verifies_password(self) -> None:
+        archive = _encrypted_zip_bytes("correct horse")
+
+        missing = inspect_archive(archive, filename="secure.zip")
+        self.assertTrue(missing.requires_password)
+        self.assertFalse(missing.password_verified)
+
+        verified = inspect_archive(
+            archive,
+            filename="secure.zip",
+            password="correct horse",
+        )
+        self.assertTrue(verified.requires_password)
+        self.assertTrue(verified.password_verified)
+
+        with self.assertRaisesRegex(ArchivePasswordError, "incorrect"):
+            inspect_archive(
+                archive,
+                filename="secure.zip",
+                password="wrong",
+            )
+
+    def test_folder_selection_extracts_only_descendants(self) -> None:
+        archive = _zip_bytes(
+            {
+                "selected/one.txt": b"one",
+                "selected/two.txt": b"two",
+                "other/three.txt": b"three",
+            }
+        )
+        stored = SimpleNamespace(
+            asset=SimpleNamespace(id="asset"),
+            version=SimpleNamespace(id="version"),
+            blob=SimpleNamespace(id="blob"),
+        )
+        with patch(
+            "govoplan_files.backend.storage.archives.create_file_asset",
+            return_value=stored,
+        ) as create_asset:
+            result = extract_archive_upload(
+                object(),  # type: ignore[arg-type]
+                tenant_id="tenant-1",
+                owner_type="user",
+                owner_id="user-1",
+                user_id="user-1",
+                archive_data=archive,
+                filename="selection.zip",
+                folder="imports",
+                campaign_id=None,
+                selected_paths=["selected"],
+            )
+
+        self.assertEqual(2, len(result))
+        self.assertEqual(
+            {"imports/selected/one.txt", "imports/selected/two.txt"},
+            {
+                call.kwargs["display_path"]
+                for call in create_asset.call_args_list
+            },
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()