"""Vulnerability mixin for ICMP/ICMPv6 findings.

Analyzes captured ICMP traffic patterns and stores vulnerability evidence.
"""
from __future__ import annotations

import csv
from collections import defaultdict

from ptnetinspector.utils.ip_utils import (
    has_additional_data,
    is_broadcast_ipv4,
    is_multicast_ipv4,
    is_valid_ipv6,
)


class ICMPMixin:
    def _extract_icmpv6_redirect_info(self, packet: str) -> dict:
        """
        Extract Target Address and Redirected Header source IP from ICMPv6 Redirect.
        
        From ICMPv6 Redirect Message:
        - tgt = target address (better first-hop router - the MitM target)
        - dst = destination address (what traffic is being redirected)
        
        From Redirected Header option embedded IPv6:
        - src = original source IP of the traffic being redirected (matches device)
        
        Returns: {
            'target_addr': str or None,
            'redirect_src_ip': str or None,
            'is_valid': bool
        }
        """
        info = {
            'target_addr': None,
            'redirect_src_ip': None,
            'is_valid': False
        }
        
        if "ICMPv6 Neighbor Discovery - Redirect" not in packet:
            return info
        
        # Extract Target Address (better first-hop router)
        # Look for 'tgt =' and extract just the IP, not everything after it
        import re
        tgt_match = re.search(r'tgt\s*=\s*([0-9a-f:]+(?:\.[0-9]+)*)', packet, re.IGNORECASE)
        if tgt_match:
            target = tgt_match.group(1).strip()
            info['target_addr'] = target

        # Extract source IP from Redirected Header option
        # The Redirected Header option contains the original packet being redirected
        # Look for '| src =' which indicates the inner IPv6 source (nested with | prefix)
        src_match = re.search(r'\|\s*src\s*=\s*([0-9a-f:]+(?:\.[0-9]+)*)', packet)
        if src_match:
            src = src_match.group(1).strip()
            # Make sure this src is inside a Redirected Header section (check context)
            # by finding Redirected Header BEFORE this src match
            redirect_header_pos = packet.rfind('Redirected Header', 0, src_match.start())
            if redirect_header_pos >= 0:
                info['redirect_src_ip'] = src
                info['is_valid'] = True

        return info
    
    def _extract_icmp_redirect_info(self, packet: str) -> dict:
        """
        Extract Gateway Address from ICMP Redirect.
        
        Returns: {
            'gateway_addr': str or None,
            'is_valid': bool
        }
        """
        info = {
            'gateway_addr': None,
            'is_valid': False
        }
        
        if "type = redirect" not in packet:
            return info
        
        # Extract Gateway Internet Address (the redirect target)
        for line in packet.split('\n'):
            if 'gw =' in line:
                parts = line.split('gw =')
                if len(parts) > 1:
                    gw = parts[1].strip()
                    if '(' in gw:
                        gw = gw.split('(')[0].strip()
                    info['gateway_addr'] = gw
                    info['is_valid'] = True
                    break
        
        return info
    
    def store_icmp_vulnerability(self) -> None:
        role_dict = self.read_role_dict(self.role_file)
        if not role_dict:
            return

        if not self.time_all_file or not has_additional_data(self.time_all_file):
            return

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

        ra_prefixes = self._load_ra_prefixes()
        all_prefixes = set()
        if self.network:
            all_prefixes.add(f"{self.network}/{self.prefix_len}")
        all_prefixes.update(ra_prefixes)

        existing_rows = self._load_existing_rows()
        new_rows = []

        icmp6_response_types = [
            "Destination unreachable",
            "Packet too big",
            "Time exceeded",
            "Parameter problem",
            "Echo Reply",
            "Redirect Message",
            "ICMP Node Information Query",
            "ICMP Node Information Response",
            "Inverse Neighbor Discovery Solicitation Message",
            "Inverse Neighbor Discovery Advertisement Message",
            "Private Experimentation",
        ]

        vuln_data = defaultdict(
            lambda: {
                "query_found": False,
                "first_idx": -1,
                "response_macs": set(),
                "response_types": defaultdict(list),
                "detected_macs": set(),
            }
        )

        vuln_checks = {}

        if getattr(self.ipver, "ipv6", False):
            vuln_checks.update(
                {
                    "icmp6_multiecho": {
                        "query": lambda p, m: "ICMPv6 Echo Request" in p and m == self.smac,
                        "response": lambda p: any(resp in p for resp in icmp6_response_types),
                        "mode": mode_general,
                        "ipver": "6",
                    },
                    "icmp6_invpacket": {
                        "query": lambda p, m: "IPv6 Extension Header - Destination Options Header" in p
                        and "Unknown Option" in p
                        and "Raw" in p
                        and m == self.smac,
                        "response": lambda p: any(resp in p for resp in icmp6_response_types),
                        "mode": mode_general,
                        "ipver": "6",
                    },
                    "icmp6_invechodo": {
                        "query": lambda p, m: "ICMPv6 Echo Request" in p
                        and "Unknown Option" in p
                        and "IPv6 Extension Header - Destination Options Header" in p
                        and m == self.smac,
                        "response": lambda p: any(resp in p for resp in icmp6_response_types),
                        "mode": mode_general,
                        "ipver": "6",
                    },
                    "icmp6_invechohbh": {
                        "query": lambda p, m: "ICMPv6 Echo Request" in p
                        and "Unknown Option" in p
                        and "IPv6 Extension Header - Hop-by-Hop Options Header" in p
                        and m == self.smac,
                        "response": lambda p: any(resp in p for resp in icmp6_response_types),
                        "mode": mode_general,
                        "ipver": "6",
                    },
                    # icmp6_outrange handled separately per new logic
                }
            )
            # Redirect detection handled later with richer logic (no simple detect here)

        if getattr(self.ipver, "ipv4", False):
            vuln_checks.update(
                {
                    "icmp_multiecho": {
                        "query": lambda p, m: "type = echo-request" in p
                        and is_multicast_ipv4(self._extract_address(p, "dst"))
                        and m == self.smac,
                        "response": lambda p: "type = echo-reply" in p,
                        "mode": mode_general,
                        "ipver": "4",
                    },
                    "icmp_brcastecho": {
                        "query": lambda p, m: "type = echo-request" in p
                        and is_broadcast_ipv4(self._extract_address(p, "dst"))
                        and m == self.smac,
                        "response": lambda p: "type = echo-reply" in p,
                        "mode": mode_general,
                        "ipver": "4",
                    },
                    "icmp_invpacket": {
                        "query": lambda p, m: "icmp echo request" in p and "type = 255" in p and m == self.smac,
                        "response": lambda p: "type = dest-unreach code = host-unreachable" in p,
                        "mode": mode_general,
                        "ipver": "4",
                    },
                    "icmp_redirect": {
                        "detect": lambda p: "type = redirect" in p,
                        "mode": mode_general,
                        "ipver": "4",
                    },
                }
            )

        # Data for new icmp6_outrange logic
        outrange_ips: set[str] = set()
        response_macs_outrange: set[str] = set()
        outrange_seen = False

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

                    # Build outrange IP set and responses per new logic
                    is_v6 = "version = 6" in packet
                    if is_v6:
                        if self._check_icmp6_outrange(packet, all_prefixes):
                            src_ip = self._extract_address(packet, "src")
                            if src_ip:
                                outrange_ips.add(src_ip)
                                outrange_seen = True

                        if outrange_ips:
                            dst_ip = self._extract_address(packet, "dst")
                            if dst_ip and dst_ip in outrange_ips:
                                response_macs_outrange.add(mac)

                    for vuln_name, vuln_check in vuln_checks.items():
                        ipver = vuln_check["ipver"]
                        version_match = f"version = {ipver}" in packet
                        if not version_match:
                            continue

                        if "detect" in vuln_check:
                            if vuln_check["detect"](packet):
                                vuln_data[vuln_name]["detected_macs"].add(mac)
                            continue

                        if not vuln_data[vuln_name]["query_found"] and vuln_check["query"](packet, mac):
                            vuln_data[vuln_name]["query_found"] = True
                            vuln_data[vuln_name]["first_idx"] = current_index

                        if vuln_data[vuln_name]["query_found"] and current_index > vuln_data[vuln_name]["first_idx"]:
                            if vuln_check["response"](packet):
                                vuln_data[vuln_name]["response_macs"].add(mac)
                                if ipver == "6":
                                    for resp_type in icmp6_response_types:
                                        if resp_type in packet and resp_type not in vuln_data[vuln_name]["response_types"][mac]:
                                            vuln_data[vuln_name]["response_types"][mac].append(resp_type)

                    current_index += 1

        except (FileNotFoundError, KeyError):
            return

        vuln_codes = {
            "icmp6_redirect": (
                "PTV-NET-MITM-6-ICMPREDIRDEV",
                "Device can be steered to forward traffic through the attacker via ICMPv6 Redirect",
                "PTV-NET-MITM-6-ICMPREDIR",
                "Network permits ICMPv6 Redirect messages that steer devices through the attacker",
            ),
            "icmp6_multiecho": (
                "PTV-NET-IDENT-ICMP-6-MULTIECHODEV",
                "Device responds to standard multicast ICMPv6 Echo Request",
                "PTV-NET-IDENT-ICMP-6-MULTIECHO",
                "Network allows the delivery of standard multicast ICMPv6 Echo Request",
            ),
            "icmp6_invechodo": (
                "PTV-NET-IDENT-ICMP-6-INVECHODODEV",
                "Device responds to invalid Destionation Option ICMPv6 Echo Request",
                "PTV-NET-IDENT-ICMP-6-INVECHODO",
                "Network allows the delivery of invalid Destionation Option ICMPv6 Echo Request",
            ),
            "icmp6_invechohbh": (
                "PTV-NET-IDENT-ICMP-6-INVECHOHBHDEV",
                "Device responds to invalid Hop-by-Hop ICMPv6 Echo Request",
                "PTV-NET-IDENT-ICMP-6-INVECHOHBH",
                "Network allows the delivery of invalid Hop-by-Hop ICMPv6 Echo Request",
            ),
            "icmp6_invpacket": (
                "PTV-NET-IDENT-ICMP-6-INVPACKETDEV",
                "Device responds to invalid ICMPv6 packets",
                "PTV-NET-IDENT-ICMP-6-INVPACKET",
                "Network allows the delivery of invalid ICMPv6 packets",
            ),
            "icmp6_outrange": (
                "PTV-NET-IDENT-ICMP-6-OUTRANGEDEV",
                "Device responds to valid or invalid ICMPv6 packets from sources outside the scanned local network range",
                "PTV-NET-IDENT-ICMP-6-OUTRANGE",
                "Network allows valid or invalid ICMPv6 packets to be delivered from sources outside the scanned local network range",
            ),
            "icmp_redirect": (
                "PTV-NET-MITM-4-ICMPREDIRDEV",
                "Device communication can be redirected using ICMP Redirect",
                "PTV-NET-MITM-4-ICMPREDIR",
                "Network does not block ICMP Redirect messages",
            ),
            "icmp_multiecho": (
                "PTV-NET-IDENT-ICMP-4-MULTIECHODEV",
                "Device responds to standard multicast ICMP Echo Request",
                "PTV-NET-IDENT-ICMP-4-MULTIECHO",
                "Network allows the delivery of standard multicast ICMP Echo Request",
            ),
            "icmp_brcastecho": (
                "PTV-NET-IDENT-ICMP-4-BRCASTECHODEV",
                "Device responds to standard broadcast ICMP Echo Request",
                "PTV-NET-IDENT-ICMP-4-BRCASTECHO",
                "Network allows the delivery of standard broadcast ICMP Echo Request",
            ),
            "icmp_invpacket": (
                "PTV-NET-IDENT-ICMP-4-INVPACKETDEV",
                "Device responds to invalid ICMP packets",
                "PTV-NET-IDENT-ICMP-4-INVPACKET",
                "Network allows the delivery of invalid ICMP packets",
            ),
        }

        for vuln_name, data in vuln_data.items():
            if vuln_name not in vuln_codes:
                continue

            vuln_mode = vuln_checks[vuln_name]["mode"]
            ipver = vuln_checks[vuln_name]["ipver"]
            dev_code, dev_desc, net_code, net_desc = vuln_codes[vuln_name]

            if "detect" in vuln_checks[vuln_name]:
                for device_id, mac in role_dict.items():
                    label = "1" if mac in data["detected_macs"] else "0"
                    row = (str(device_id), mac, vuln_mode, ipver, dev_code, dev_desc, label)
                    self._add_vuln_if_new(row, existing_rows, new_rows)

                net_label = "1" if data["detected_macs"] else "0"
                net_row = ("Network", "", vuln_mode, ipver, net_code, net_desc, net_label)
                self._add_vuln_if_new(net_row, existing_rows, new_rows)
                continue

            if not data["query_found"]:
                for device_id, mac in role_dict.items():
                    row = (str(device_id), mac, vuln_mode, ipver, dev_code, dev_desc, "2")
                    self._add_vuln_if_new(row, existing_rows, new_rows)

                net_row = ("Network", "", vuln_mode, ipver, 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():
                    desc = dev_desc
                    if ipver == "6" and mac in data["response_macs"]:
                        desc = f"{dev_desc} with ICMPv6 message"

                    label = "1" if mac in data["response_macs"] else "0"
                    row = (str(device_id), mac, vuln_mode, ipver, dev_code, desc, label)
                    self._add_vuln_if_new(row, existing_rows, new_rows)

                net_label = "1" if data["response_macs"] else "0"
                net_row = ("Network", "", vuln_mode, ipver, net_code, net_desc, net_label)
                self._add_vuln_if_new(net_row, existing_rows, new_rows)

        self._write_vulnerabilities(new_rows)

        # Emit rows for icmp6_outrange per new logic
        existing_rows = self._load_existing_rows()
        new_outrange_rows = []

        dev_code, dev_desc, net_code, net_desc = vuln_codes["icmp6_outrange"]
        vuln_mode = mode_general
        ipver_value = "6"

        role_macs = set(role_dict.values())
        reacting_macs = {m for m in response_macs_outrange if m in role_macs}

        if not outrange_seen:
            for device_id, mac in role_dict.items():
                row = (str(device_id), mac, vuln_mode, ipver_value, dev_code, dev_desc, "2")
                self._add_vuln_if_new(row, existing_rows, new_outrange_rows)

            net_row = ("Network", "", vuln_mode, ipver_value, net_code, net_desc, "2")
            self._add_vuln_if_new(net_row, existing_rows, new_outrange_rows)
        else:
            for device_id, mac in role_dict.items():
                label = "1" if mac in reacting_macs else "0"
                row = (str(device_id), mac, vuln_mode, ipver_value, dev_code, dev_desc, label)
                self._add_vuln_if_new(row, existing_rows, new_outrange_rows)

            net_label = "1" if reacting_macs else "0"
            net_row = ("Network", "", vuln_mode, ipver_value, net_code, net_desc, net_label)
            self._add_vuln_if_new(net_row, existing_rows, new_outrange_rows)

        if new_outrange_rows:
            self._write_vulnerabilities(new_outrange_rows)

        existing_rows = self._load_existing_rows()
        new_redirect_rows = []

        redirect_vulns = {
            "4": (
                "PTV-NET-MITM-4-ICMPREDIRDEV",
                "Device communication can be redirected using ICMP Redirect",
                "PTV-NET-MITM-4-ICMPREDIR",
                "Network does not block ICMP Redirect messages",
                mode_general,
            ),
            "6": (
                "PTV-NET-MITM-6-ICMPREDIRDEV",
                "Device can be steered to forward traffic through the attacker via ICMPv6 Redirect",
                "PTV-NET-MITM-6-ICMPREDIR",
                "Network permits ICMPv6 Redirect messages that steer devices through the attacker",
                mode_redirect_v6,
            ),
        }

        ipver_to_process = []
        if getattr(self.ipver, "ipv4", False):
            ipver_to_process.append("4")
        if getattr(self.ipver, "ipv6", False) and mode_redirect_v6:
            ipver_to_process.append("6")

        for ipver_value in ipver_to_process:
            if ipver_value not in redirect_vulns:
                continue

            dev_code, dev_desc, net_code, net_desc, vuln_mode = redirect_vulns[ipver_value]

            # Track redirects with MitM criteria vs others
            vuln_devices = set()  # MACs of devices whose traffic is redirected to SCANNER (MitM)
            redirect_devices = set()  # MACs of devices with redirects to OTHER targets (not scanner)
            any_redirect_found = False
            any_mitm_redirect_found = False
            
            # Build IP-to-MAC mapping from addresses_unfiltered.csv
            from ptnetinspector.utils.path import get_csv_path
            addresses_file = get_csv_path("addresses_unfiltered.csv", self.interface)
            ip_to_device_mac = {}
            
            try:
                with open(addresses_file, "r", newline="") as csvfile:
                    reader = csv.DictReader(csvfile)
                    for row in reader:
                        ip = row.get("IP", "").strip()
                        mac = row.get("MAC", "").strip()
                        if ip and mac:
                            ip_to_device_mac[ip] = mac
            except (FileNotFoundError, KeyError):
                pass
            
            try:
                with open(self.time_all_file, "r", newline="") as csvfile:
                    reader = csv.DictReader(csvfile)
                    for row in reader:
                        packet = row.get("packet", "")
                        
                        if ipver_value == "6":
                            redirect_info = self._extract_icmpv6_redirect_info(packet)
                            if redirect_info['is_valid']:
                                any_redirect_found = True
                                
                                # Only consider vulnerability if we can extract source IP
                                if not redirect_info['redirect_src_ip']:
                                    # Can't extract source from Redirected Header option
                                    # Skip this redirect - can't assess device involvement
                                    continue
                                
                                # Check if source IP matches a known device
                                src_ip = redirect_info['redirect_src_ip']
                                if src_ip not in ip_to_device_mac:
                                    # Source doesn't match any known device, skip
                                    continue
                                
                                device_mac = ip_to_device_mac[src_ip]
                                
                                # Now check target address
                                if redirect_info['target_addr'] and self.sip:
                                    sip_list = self.sip if isinstance(self.sip, list) else [self.sip]
                                    if redirect_info['target_addr'] in sip_list:
                                        # Target matches scanner - this is a MitM attack
                                        vuln_devices.add(device_mac)
                                        any_mitm_redirect_found = True
                                    else:
                                        # Target does NOT match scanner - redirect to other device
                                        redirect_devices.add(device_mac)
                        
                        elif ipver_value == "4":
                            redirect_info = self._extract_icmp_redirect_info(packet)
                            if redirect_info['is_valid']:
                                any_redirect_found = True
                                # For IPv4, similar logic needed
                                # Would need source IP extraction from ICMP redirect message
                                pass
            except (FileNotFoundError, KeyError):
                pass
            
            # Assign labels. Always recompute: drop existing redirect rows for this code/mode/ipver
            # Remove any prior redirect rows (both device and network codes) for this mode/ipver
            existing_rows = {
                r
                for r in existing_rows
                if not (
                    len(r) >= 5
                    and r[3] == ipver_value
                    and r[2] == vuln_mode
                    and r[4] in {dev_code, net_code}
                )
            }

            if not any_redirect_found:
                # No redirects detected at all - label all devices as 2 (N/A)
                for device_id, mac in role_dict.items():
                    row = (str(device_id), mac, vuln_mode, ipver_value, dev_code, dev_desc, "2")
                    self._add_vuln_if_new(row, existing_rows, new_redirect_rows)
                net_row = ("Network", "", vuln_mode, ipver_value, net_code, net_desc, "2")
                self._add_vuln_if_new(net_row, existing_rows, new_redirect_rows)
            else:
                # Redirects found - label devices based on:
                # 1 = vulnerable (traffic redirected to scanner - MitM)
                # 0 = redirected but not to scanner (not a MitM against scanner)
                # 2 = N/A (no source IP match or couldn't extract)
                for device_id, mac in role_dict.items():
                    if mac in vuln_devices:
                        label = "1"  # MitM vulnerability
                    elif mac in redirect_devices:
                        label = "0"  # Redirect exists but not to scanner
                    else:
                        label = "2"  # N/A - not involved in any valid redirect
                    row = (str(device_id), mac, vuln_mode, ipver_value, dev_code, dev_desc, label)
                    self._add_vuln_if_new(row, existing_rows, new_redirect_rows)
                
                # Network label tied to device assessment:
                # - If any device is labelled 1 (MitM) or 0 (redirect observed), network is vulnerable (1)
                # - If all devices are N/A (2), network is N/A (2)
                if vuln_devices:
                    net_label = "1"
                elif redirect_devices:
                    net_label = "0"
                else:
                    net_label = "2"
                net_row = ("Network", "", vuln_mode, ipver_value, net_code, net_desc, net_label)
                self._add_vuln_if_new(net_row, existing_rows, new_redirect_rows)

        # Rewrite vulnerability file for redirect rows: remove old entries of this vuln/mode/ipver, then add new ones
        try:
            fieldnames = ["ID", "MAC", "Mode", "IPver", "Code", "Description", "Label"]
            # Start from filtered existing_rows (tuples) and convert to dicts
            remaining_rows = [dict(zip(fieldnames, r)) for r in existing_rows]
            combined_rows = remaining_rows + new_redirect_rows

            # Always rewrite file to avoid stale redirect labels
            import os
            os.makedirs(os.path.dirname(self.vulnerability_file), exist_ok=True)
            with open(self.vulnerability_file, "w", newline="") as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(combined_rows)
        except Exception:
            pass
