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