Initial release: offline NetScaler triage with Sigma rules and MIT license

tests/test_preparation.py

@@ -0,0 +1,173 @@
+import gzip
+import io
+import tempfile
+import unittest
+import zipfile
+from pathlib import Path
+from typing import IO, Any, cast
+from unittest.mock import patch
+
+from netscaler_triage.ingest import InputError, iter_records
+from netscaler_triage.preparation import prepare_inputs
+
+
+class PreparationTests(unittest.TestCase):
+    def test_corrupt_archive_does_not_abort_other_inputs(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp).resolve()
+            (root / "bad.zip").write_bytes(b"not a zip")
+            (root / "good.log").write_text("evidence\n")
+            groups, manifest = prepare_inputs([root / "bad.zip", root / "good.log"], root / "out")
+            self.assertEqual(len(groups), 1)
+            errors = [row for row in manifest if row["action"] == "error"]
+            self.assertEqual(errors[0]["error"], "ArchiveError")
+            self.assertIn("sha256", errors[0])
+
+    def test_unreadable_file_is_manifest_error(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp).resolve()
+            source = root / "source"
+            source.mkdir()
+            (source / "bad").write_text("private")
+            (source / "good").write_text("evidence")
+            original = Path.open
+
+            def guarded(path: Path, *args: Any, **kwargs: Any) -> IO[Any]:
+                if path == source / "bad":
+                    raise PermissionError(13, "denied", str(path))
+                return cast(IO[Any], original(path, *args, **kwargs))
+
+            with patch.object(Path, "open", guarded):
+                groups, manifest = prepare_inputs([source], root / "out")
+            self.assertEqual(len(groups), 1)
+            self.assertEqual([r["source"] for r in manifest if r["action"] == "error"], ["bad"])
+            self.assertEqual(
+                [p.name for _, folder in groups for p in folder.rglob("*") if p.is_file()], ["good"]
+            )
+
+    def test_password_and_combined_archive_budget(self) -> None:
+        import py7zr
+
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp).resolve()
+            encrypted = root / "private.7z"
+            with py7zr.SevenZipFile(
+                encrypted, "w", password="test-secret", header_encryption=True
+            ) as z:
+                z.writestr(b"log evidence\n", "log/messages")
+            groups, manifest = prepare_inputs([encrypted], root / "good", password="test-secret")
+            self.assertEqual(len(groups), 1)
+            self.assertNotIn("test-secret", str(manifest))
+            archives = []
+            for name in ("a.zip", "b.zip"):
+                path = root / name
+                with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_DEFLATED) as z:
+                    z.writestr("messages", "A" * 1000)
+                archives.append(path)
+            limit = sum(p.stat().st_size for p in archives) + 1999
+            with self.assertRaises(InputError):
+                prepare_inputs(archives, root / "limited", max_bytes=limit)
+            self.assertFalse((root / "limited").exists())
+
+    def test_multiple_appliances_and_gzip(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp).resolve()
+            source = root / "export.zip"
+            with zipfile.ZipFile(source, "w") as archive:
+                archive.writestr("wrapper/QS/var/log/messages", b"first\n")
+                archive.writestr("wrapper/PROD/var/log/messages.gz", gzip.compress(b"second\n"))
+            before = source.read_bytes()
+            groups, manifest = prepare_inputs([source], root / "prepared")
+            self.assertEqual(len(groups), 2)
+            records = [r for name, folder in groups for r in iter_records(folder, name)]
+            self.assertEqual({r["raw"] for r in records}, {"first", "second"})
+            self.assertFalse(any(r.get("error") for r in records))
+            self.assertEqual(before, source.read_bytes())
+            self.assertEqual(sum(r["action"] == "scan_root" for r in manifest), 2)
+
+    def test_toolkit_exclusions_preserve_script_evidence(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp).resolve()
+            source = root / "source"
+            source.mkdir()
+            (source / "pyproject.toml").write_text('[project]\nname = "netscaler-triage"\n')
+            (source / "src").mkdir()
+            (source / "src/tool.py").write_text("tool")
+            (source / "results").mkdir()
+            (source / "results/old.log").write_text("old")
+            (source / "messages").write_text("evidence\n")
+            (source / "malware.py").write_text('print("evidence")\n')
+            groups, manifest = prepare_inputs([source], source / "new-output")
+            files = {p.name for _, folder in groups for p in folder.rglob("*") if p.is_file()}
+            self.assertEqual(files, {"messages", "malware.py"})
+            self.assertEqual(
+                {r["source"] for r in manifest if r["action"] == "excluded"},
+                {"src", "results", "pyproject.toml", "new-output"},
+            )
+
+    def test_plain_evidence_src_is_not_toolkit(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp).resolve()
+            source = root / "source"
+            (source / "src").mkdir(parents=True)
+            (source / "src/malware.py").write_text("code evidence")
+            groups, _ = prepare_inputs([source], root / "out")
+            self.assertEqual(sum(len(list(p.rglob("malware.py"))) for _, p in groups), 1)
+
+    def test_gzip_wrapped_zip_and_nested_named_archive(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp).resolve()
+            source = root / "source.zip"
+            zipped = io.BytesIO()
+            with zipfile.ZipFile(zipped, "w") as archive:
+                archive.writestr("log/messages", "hello\n")
+            with zipfile.ZipFile(source, "w") as archive:
+                archive.writestr("inner.gz", gzip.compress(zipped.getvalue()))
+            groups, _ = prepare_inputs([source], root / "out")
+            records = [r for name, folder in groups for r in iter_records(folder, name)]
+            self.assertEqual([r["raw"] for r in records], ["hello"])
+
+    def test_combined_quota_and_nooverwrite(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp).resolve()
+            a = root / "a.log"
+            b = root / "b.log"
+            a.write_bytes(b"a" * 6)
+            b.write_bytes(b"b" * 6)
+            with self.assertRaises(InputError):
+                prepare_inputs([a, b], root / "out", max_bytes=10)
+            self.assertFalse((root / "out").exists())
+            (root / "out").mkdir()
+            (root / "out/keep").write_text("safe")
+            with self.assertRaises(FileExistsError):
+                prepare_inputs([a], root / "out")
+            self.assertEqual((root / "out/keep").read_text(), "safe")
+
+    def test_symlink_rejected_and_source_preserved(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp).resolve()
+            source = root / "source"
+            source.mkdir()
+            (source / "real").write_text("evidence")
+            (source / "link").symlink_to("real")
+            with self.assertRaises(InputError):
+                prepare_inputs([source], root / "out")
+            self.assertFalse((root / "out").exists())
+            self.assertEqual((source / "real").read_text(), "evidence")
+
+    def test_two_same_named_exports_stay_distinct(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp).resolve()
+            sources = []
+            for n in ("one", "two"):
+                source = root / n / "log"
+                source.mkdir(parents=True)
+                (source / "messages").write_text(n)
+                sources.append(source)
+            groups, _ = prepare_inputs(sources, root / "out")
+            self.assertEqual(len(groups), 2)
+            self.assertEqual(len({name for name, _ in groups}), 2)
+
+
+if __name__ == "__main__":
+    unittest.main()