"""Vulnerability mixin for mDNS-related issues.

Detects insecure mDNS behaviors based on captured traffic patterns.
"""
from __future__ import annotations

import csv

from ptnetinspector.utils.ip_utils import has_additional_data
from ptnetinspector.vulnerability.dns_helpers import (
    contains_dns_qtype,
    contains_dns_type,
    dns_record_types_no_ptr,
)


class MDNSMixin:
    def store_MDNS_vulnerability(self) -> None:
        role_dict = self.read_role_dict(self.role_file)
        if not role_dict:
            return

        ipver_values = self._get_ipver_values()
        if not ipver_values:
            return

        existing_rows = self._load_existing_rows()
        new_rows = []

        for ipver_value in ipver_values:
            self._process_mdns_dnssd(role_dict, ipver_value, existing_rows, new_rows)

        self._write_vulnerabilities(new_rows)

    def _process_mdns_dnssd(self, role_dict, ipver_value: str, existing_rows, new_rows) -> None:
        mdns_std_query_found = False
        mdns_ptr_query_found = False
        dnssd_query_found = False
        first_mdns_std_query_index = -1
        first_mdns_ptr_query_index = -1
        first_dnssd_query_index = -1
        mdns_std_response_macs = set()
        mdns_ptr_response_macs = set()
        dnssd_response_macs = set()

        non_ptr_types = dns_record_types_no_ptr()

        try:
            with open(self.time_all_file, "r", newline="") as csvfile:
                reader = csv.DictReader(csvfile)
                current_index = 0

                for row in reader:
                    packet = row.get("packet", "")
                    mac = row.get("MAC", "").strip()

                    version_match = (ipver_value == "4" and "version = 4" in packet) or (
                        ipver_value == "6" and "version = 6" in packet
                    )

                    if not version_match:
                        current_index += 1
                        continue

                    is_mdns_query = "dport = mdns" in packet and "qr = 0" in packet and mac == self.smac

                    if is_mdns_query:
                        if "qtype = PTR" in packet and not mdns_ptr_query_found:
                            mdns_ptr_query_found = True
                            first_mdns_ptr_query_index = current_index
                        elif not mdns_std_query_found and contains_dns_qtype(packet, non_ptr_types):
                            mdns_std_query_found = True
                            first_mdns_std_query_index = current_index

                    if (
                        not dnssd_query_found
                        and "dport = mdns" in packet
                        and "qr = 0" in packet
                        and "dns-sd._udp.local" in packet
                        and mac == self.smac
                    ):
                        dnssd_query_found = True
                        first_dnssd_query_index = current_index

                    if (
                        mdns_std_query_found
                        and current_index > first_mdns_std_query_index
                        and "dport = mdns" in packet
                        and "qr = 1" in packet
                    ):
                        if contains_dns_type(packet, non_ptr_types):
                            mdns_std_response_macs.add(mac)

                    if (
                        mdns_ptr_query_found
                        and current_index > first_mdns_ptr_query_index
                        and "dport = mdns" in packet
                        and "qr = 1" in packet
                        and "type = PTR" in packet
                    ):
                        mdns_ptr_response_macs.add(mac)

                    if (
                        dnssd_query_found
                        and current_index > first_dnssd_query_index
                        and "dport = mdns" in packet
                        and "qr = 1" in packet
                        and "dns-sd._udp.local" in packet
                    ):
                        dnssd_response_macs.add(mac)

                    current_index += 1

            mode = self._filter_modes("a,a+")
            if not mode:
                return

            if ipver_value == "4":
                mdns_std_device_code = "PTV-NET-IDENT-4-MDNSDEV"
                mdns_std_device_desc = "Device responds to IPv4 MDNS standard DNS record queries"
                mdns_std_net_code = "PTV-NET-IDENT-4-MDNS"
                mdns_std_net_desc = "Network allows the delivery of IPv4 MDNS standard DNS record queries"
            else:
                mdns_std_device_code = "PTV-NET-IDENT-6-MDNSDEV"
                mdns_std_device_desc = "Device responds to IPv6 MDNS standard DNS record queries"
                mdns_std_net_code = "PTV-NET-IDENT-6-MDNS"
                mdns_std_net_desc = "Network allows the delivery of IPv6 MDNS standard DNS record queries"

            self._add_protocol_vulnerabilities(
                role_dict,
                mdns_std_query_found,
                mdns_std_response_macs,
                ipver_value,
                mdns_std_device_code,
                mdns_std_device_desc,
                mdns_std_net_code,
                mdns_std_net_desc,
                mode,
                existing_rows,
                new_rows,
            )

            if ipver_value == "4":
                mdns_ptr_device_code = "PTV-NET-IDENT-4-MDNS-PTRDEV"
                mdns_ptr_device_desc = "Device responds to IPv4 MDNS PTR packets"
                mdns_ptr_net_code = "PTV-NET-IDENT-4-MDNS-PTR"
                mdns_ptr_net_desc = "Network allows the delivery of IPv4 MDNS PTR packets"
            else:
                mdns_ptr_device_code = "PTV-NET-IDENT-6-MDNS-PTRDEV"
                mdns_ptr_device_desc = "Device responds to IPv6 MDNS PTR packets"
                mdns_ptr_net_code = "PTV-NET-IDENT-6-MDNS-PTR"
                mdns_ptr_net_desc = "Network allows the delivery of IPv6 MDNS PTR packets"

            self._add_protocol_vulnerabilities(
                role_dict,
                mdns_ptr_query_found,
                mdns_ptr_response_macs,
                ipver_value,
                mdns_ptr_device_code,
                mdns_ptr_device_desc,
                mdns_ptr_net_code,
                mdns_ptr_net_desc,
                mode,
                existing_rows,
                new_rows,
            )

            if ipver_value == "4":
                dnssd_device_code = "PTV-NET-IDENT-4-DNSSDDEV"
                dnssd_device_desc = "Device responds to IPv4 DNS Service Discovery packets"
                dnssd_net_code = "PTV-NET-IDENT-4-DNSSD"
                dnssd_net_desc = "Network allows devices discovery with IPv4 DNS Service Discovery"
            else:
                dnssd_device_code = "PTV-NET-IDENT-6-DNSSDDEV"
                dnssd_device_desc = "Device responds to IPv6 DNS Service Discovery packets"
                dnssd_net_code = "PTV-NET-IDENT-6-DNSSD"
                dnssd_net_desc = "Network allows devices discovery with IPv6 DNS Service Discovery"

            self._add_protocol_vulnerabilities(
                role_dict,
                dnssd_query_found,
                dnssd_response_macs,
                ipver_value,
                dnssd_device_code,
                dnssd_device_desc,
                dnssd_net_code,
                dnssd_net_desc,
                mode,
                existing_rows,
                new_rows,
            )

        except (FileNotFoundError, KeyError):
            pass

    def _add_protocol_vulnerabilities(
        self,
        role_dict,
        query_found: bool,
        response_macs: set[str],
        ipver_value: str,
        device_code: str,
        device_desc: str,
        net_code: str,
        net_desc: str,
        mode: str,
        existing_rows,
        new_rows,
    ) -> None:
        if not query_found:
            for device_id, mac in role_dict.items():
                row = (str(device_id), mac, mode, ipver_value, device_code, device_desc, "2")
                self._add_vuln_if_new(row, existing_rows, new_rows)

            net_row = ("Network", "", mode, ipver_value, net_code, net_desc, "2")
            self._add_vuln_if_new(net_row, existing_rows, new_rows)
        else:
            for device_id, mac in role_dict.items():
                label = "1" if mac in response_macs else "0"
                row = (str(device_id), mac, mode, ipver_value, device_code, device_desc, label)
                self._add_vuln_if_new(row, existing_rows, new_rows)

            net_label = "1" if response_macs else "0"
            net_row = ("Network", "", mode, ipver_value, net_code, net_desc, net_label)
            self._add_vuln_if_new(net_row, existing_rows, new_rows)
