63 lines
2.8 KiB
Python
63 lines
2.8 KiB
Python
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import py7zr
|
|
import pyzipper
|
|
|
|
from netscaler_triage.ingest import InputError, extract_archive
|
|
|
|
|
|
class ArchivePasswordTests(unittest.TestCase):
|
|
def test_encrypted_7z_and_aes_zip(self) -> None:
|
|
password = "sensitive-Ü-test-password"
|
|
with tempfile.TemporaryDirectory() as folder:
|
|
root = Path(folder).resolve()
|
|
for extension in ("7z", "zip"):
|
|
archive = root / ("input." + extension)
|
|
if extension == "7z":
|
|
with py7zr.SevenZipFile(
|
|
archive, "w", password=password, header_encryption=True
|
|
) as z:
|
|
z.writestr(b"log data\n", "logs/access.log")
|
|
else:
|
|
with pyzipper.AESZipFile(archive, "w", encryption=pyzipper.WZ_AES) as z:
|
|
z.setpassword(password.encode())
|
|
z.writestr("logs/access.log", b"log data\n")
|
|
original = archive.read_bytes()
|
|
destination = root / (extension + "-correct")
|
|
manifest = extract_archive(archive, destination, password=password)
|
|
self.assertEqual((destination / "logs/access.log").read_bytes(), b"log data\n")
|
|
self.assertNotIn(password, json.dumps(manifest))
|
|
for wrong in (None, "other-secret"):
|
|
failed = root / (extension + "-failed")
|
|
with self.assertRaises(InputError) as caught:
|
|
extract_archive(archive, failed, password=wrong)
|
|
self.assertNotIn(password, str(caught.exception))
|
|
self.assertNotIn("other-secret", str(caught.exception))
|
|
self.assertFalse(failed.exists())
|
|
self.assertEqual(archive.read_bytes(), original)
|
|
|
|
def test_nested_encryption_and_quota(self) -> None:
|
|
with tempfile.TemporaryDirectory() as folder:
|
|
root = Path(folder).resolve()
|
|
inner = root / "nested.7z"
|
|
with py7zr.SevenZipFile(inner, "w", password="secret", header_encryption=True) as z:
|
|
z.writestr(b"nested log", "messages")
|
|
outer = root / "outer.zip"
|
|
with pyzipper.AESZipFile(outer, "w", encryption=pyzipper.WZ_AES) as z:
|
|
z.setpassword(b"secret")
|
|
z.write(inner, "nested.7z")
|
|
target = root / "output"
|
|
extract_archive(outer, target, password="secret")
|
|
self.assertEqual((target / "nested.7z.extracted/messages").read_bytes(), b"nested log")
|
|
failed = root / "quota"
|
|
with self.assertRaisesRegex(InputError, "size limit"):
|
|
extract_archive(outer, failed, password="secret", max_bytes=2)
|
|
self.assertFalse(failed.exists())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|