"""Tests for vulnerability detection logic changes (protocol codes and RA ICMP6-OUTRANGE)."""
import csv
import tempfile
from pathlib import Path
from unittest.mock import patch

from ptnetinspector.vulnerability import Vulnerability


class SimpleIPVer:
    """Lightweight IP version holder for tests."""

    def __init__(self, ipv4: bool = True, ipv6: bool = True):
        self.ipv4 = ipv4
        self.ipv6 = ipv6


def _write_role_file(path: Path, entries: list[tuple[int, str]]):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", newline="") as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=["MAC", "Device_Number", "Role"])
        writer.writeheader()
        for dev_num, mac in entries:
            writer.writerow({"MAC": mac, "Device_Number": dev_num, "Role": "Host"})


def _write_time_all(path: Path, rows: list[dict[str, str]]):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", newline="") as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=["time", "MAC", "packet"])
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def _read_codes(path: Path) -> set[tuple[str, str]]:
    with open(path, newline="") as csvfile:
        reader = csv.DictReader(csvfile)
        return {(r["Code"], r["IPver"]) for r in reader}


def _read_rows_by_code(path: Path, code: str) -> list[dict[str, str]]:
    with open(path, newline="") as csvfile:
        reader = csv.DictReader(csvfile)
        return [r for r in reader if r["Code"] == code]


def test_mdns_dnssd_llmnr_codes_per_ip_version():
    """Verify MDNS/DNS-SD/LLMNR emit IP-version-specific codes."""
    with tempfile.TemporaryDirectory() as tmpdir:
        role_file = Path(tmpdir) / "role_node.csv"
        time_all = Path(tmpdir) / "time_all.csv"
        vuln_file = Path(tmpdir) / "vulnerability.csv"

        smac = "aa:aa:aa:aa:aa:aa"
        device_mac = "11:22:33:44:55:66"

        _write_role_file(role_file, [(1, device_mac)])

        rows = [
            {"time": "0", "MAC": smac, "packet": "version = 4 dport = mdns qr = 0"},
            {"time": "1", "MAC": device_mac, "packet": "version = 4 dport = mdns qr = 1"},
            {"time": "2", "MAC": smac, "packet": "version = 4 dport = mdns qr = 0 dns-sd._udp.local"},
            {"time": "3", "MAC": device_mac, "packet": "version = 4 dport = mdns qr = 1 dns-sd._udp.local"},
            {"time": "4", "MAC": smac, "packet": "version = 6 dport = mdns qr = 0"},
            {"time": "5", "MAC": device_mac, "packet": "version = 6 dport = mdns qr = 1"},
            {"time": "6", "MAC": smac, "packet": "version = 6 dport = mdns qr = 0 dns-sd._udp.local"},
            {"time": "7", "MAC": device_mac, "packet": "version = 6 dport = mdns qr = 1 dns-sd._udp.local"},
            {"time": "8", "MAC": smac, "packet": "version = 4 dport = 5355 qr = 0"},
            {"time": "9", "MAC": device_mac, "packet": "version = 4 dport = 5355 qr = 1"},
            {"time": "10", "MAC": smac, "packet": "version = 6 dport = 5355 qr = 0"},
            {"time": "11", "MAC": device_mac, "packet": "version = 6 dport = 5355 qr = 1"},
        ]
        _write_time_all(time_all, rows)

        vuln = Vulnerability(
            interface="eth0",
            mode=["a", "a+"],
            ipver=SimpleIPVer(ipv4=True, ipv6=True),
            smac=smac,
            network="",
            prefix_len=64,
            rpref=0,
            dns=[],
            role_file=str(role_file),
            time_all_file=str(time_all),
            vulnerability_file=str(vuln_file),
        )

        vuln.store_MDNS_vulnerability()
        vuln.store_llmnr_vulnerability()

        codes = _read_codes(vuln_file)
        expected = {
            ("PTV-NET-IDENT-MDNS-PTRDEV", "4"),
            ("PTV-NET-IDENT-MDNS-PTR", "4"),
            ("PTV-NET-IDENT-MDNS-PTR6DEV", "6"),
            ("PTV-NET-IDENT-MDNS-PTR6", "6"),
            ("PTV-NET-IDENT-DNS-SDDEV", "4"),
            ("PTV-NET-IDENT-DNS-SD", "4"),
            ("PTV-NET-IDENT-DNS-SD6DEV", "6"),
            ("PTV-NET-IDENT-DNS-SD6", "6"),
            ("PTV-NET-IDENT-LLMNR-PTRDEV", "4"),
            ("PTV-NET-IDENT-LLMNR-PTR", "4"),
            ("PTV-NET-IDENT-LLMNR-PTR6DEV", "6"),
            ("PTV-NET-IDENT-LLMNR-PTR6", "6"),
        }

        for code in expected:
            assert code in codes


def test_ra_icmp6_outrange_detection():
    """Ensure RA processing flags ICMPv6 out-of-range traffic for device and network."""
    with tempfile.TemporaryDirectory() as tmpdir:
        role_file = Path(tmpdir) / "role_node.csv"
        time_all = Path(tmpdir) / "time_all.csv"
        vuln_file = Path(tmpdir) / "vulnerability.csv"

        smac = "aa:aa:aa:aa:aa:aa"
        device_mac = "11:22:33:44:55:66"
        _write_role_file(role_file, [(1, device_mac)])

        rows = [
            {"time": "0", "MAC": smac, "packet": "ICMPv6 Neighbor Discovery - Router Advertisement"},
            {"time": "1", "MAC": device_mac, "packet": "ICMPv6 src = 2001:db9::1 dst = fe80::1"},
        ]
        _write_time_all(time_all, rows)

        vuln = Vulnerability(
            interface="eth0",
            mode=["a+"],
            ipver=SimpleIPVer(ipv4=False, ipv6=True),
            smac=smac,
            network="2001:db8::",
            prefix_len=64,
            rpref=0,
            dns=[],
            role_file=str(role_file),
            time_all_file=str(time_all),
            vulnerability_file=str(vuln_file),
        )

        with patch("ptnetinspector.vulnerability.Interface.get_interface_ipv6_ips", return_value=["2001:db8::123"]):
            vuln.store_ra_vulnerability()

        dev_rows = _read_rows_by_code(vuln_file, "PTV-NET-IDENT-ICMP6-OUTRANGEDEV")
        net_rows = _read_rows_by_code(vuln_file, "PTV-NET-IDENT-ICMP6-OUTRANGE")

        assert any(r["Label"] == "1" for r in dev_rows)
        assert any(r["Label"] == "1" for r in net_rows)
