import ipaddress
import os
import struct
import tempfile
import unittest

from app.forensics import ForensicPcapRing


def ipv4_frame(src: str, dst: str, payload: bytes = b"evidence") -> bytes:
    ethernet = b"\x00\x11\x22\x33\x44\x55\x66\x77\x88\x99\xaa\xbb\x08\x00"
    total = 20 + len(payload)
    header = bytearray(20)
    header[0] = 0x45
    header[2:4] = total.to_bytes(2, "big")
    header[8] = 64
    header[9] = 6
    header[12:16] = ipaddress.IPv4Address(src).packed
    header[16:20] = ipaddress.IPv4Address(dst).packed
    return ethernet + bytes(header) + payload


class ForensicPcapTests(unittest.TestCase):
    def test_blocks_mode_persists_only_when_capture_is_requested(self):
        with tempfile.TemporaryDirectory() as td:
            ring = ForensicPcapRing("blocks", td, window_seconds=60, memory_mb=1)
            ring.observe(ipv4_frame("192.0.2.10", "198.51.100.4"))
            self.assertEqual(os.listdir(td), [])

            result = ring.capture_target("198.51.100.4", label="sid-42")
            self.assertIsNotNone(result)
            path = os.path.join(td, result["name"])
            self.assertTrue(os.path.isfile(path))
            self.assertEqual(result["packet_count"], 1)
            with open(path, "rb") as handle:
                magic = struct.unpack("<I", handle.read(4))[0]
            self.assertEqual(magic, 0xA1B2C3D4)

    def test_alerts_mode_leaves_persistence_to_suricata(self):
        with tempfile.TemporaryDirectory() as td:
            ring = ForensicPcapRing("alerts", td)
            ring.observe(ipv4_frame("192.0.2.10", "198.51.100.4"))
            self.assertIsNone(ring.capture_target("198.51.100.4"))
            self.assertEqual(os.listdir(td), [])


if __name__ == "__main__":
    unittest.main()