174 lines
8.0 KiB
Python
174 lines
8.0 KiB
Python
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()
|