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()