"""
Exploit mitigation analyzer module using radare2.
Checks for security features like DEP, ASLR, SafeSEH, CFG, etc.
"""

import logging
from typing import Any

from ..utils.r2_helpers import safe_cmd, safe_cmdj
from ..utils.r2_suppress import silent_cmdj

logger = logging.getLogger(__name__)


class ExploitMitigationAnalyzer:
    """Analyzes exploit mitigation features in PE files."""

    # DllCharacteristics flags
    DLL_CHARACTERISTICS = {
        0x0020: "HIGH_ENTROPY_VA",  # ASLR with high entropy
        0x0040: "DYNAMIC_BASE",  # ASLR enabled
        0x0080: "FORCE_INTEGRITY",  # Force code integrity checks
        0x0100: "NX_COMPAT",  # DEP enabled
        0x0200: "NO_ISOLATION",  # No isolation
        0x0400: "NO_SEH",  # No SEH
        0x0800: "NO_BIND",  # Do not bind image
        0x1000: "APPCONTAINER",  # Image should run in AppContainer
        0x2000: "WDM_DRIVER",  # WDM driver
        0x4000: "GUARD_CF",  # Control Flow Guard enabled
        0x8000: "TERMINAL_SERVER_AWARE",  # Terminal server aware
    }

    def __init__(self, r2):
        """
        Initialize the Exploit Mitigation analyzer.

        Args:
            r2: Radare2 instance
        """
        self.r2 = r2

    def analyze(self) -> dict[str, Any]:
        """
        Analyze exploit mitigation features in the PE file.

        Returns:
            Dictionary containing mitigation information
        """
        try:
            result = {
                "mitigations": {},
                "dll_characteristics": {},
                "load_config": {},
                "security_score": 0,
                "recommendations": [],
                "vulnerabilities": [],
            }

            # Check DllCharacteristics flags
            self._check_dll_characteristics(result)

            # Check for SafeSEH
            self._check_safe_seh(result)

            # Check Load Configuration Directory
            self._check_load_config(result)

            # Check for stack cookies
            self._check_stack_cookies(result)

            # Check for signed binary
            self._check_authenticode(result)

            # Check PE header security features
            self._check_pe_security_features(result)

            # Calculate security score
            self._calculate_security_score(result)

            # Generate recommendations
            self._generate_recommendations(result)

            return result

        except Exception as e:
            logger.error(f"Error analyzing exploit mitigations: {e}")
            return {"mitigations": {}, "error": str(e)}

    def _check_dll_characteristics(self, result: dict[str, Any]):
        """Check DllCharacteristics flags for security features."""
        try:
            # Get optional header
            optional_header = silent_cmdj(self.r2, "iHj", {})
            if not optional_header or not isinstance(optional_header, dict):
                return

            dll_characteristics = optional_header.get("dll_characteristics", 0)

            # Parse individual flags
            flags_found = {}
            mitigations = {}

            for flag_value, flag_name in self.DLL_CHARACTERISTICS.items():
                if dll_characteristics & flag_value:
                    flags_found[flag_name] = True

            result["dll_characteristics"] = {
                "value": hex(dll_characteristics),
                "flags": flags_found,
            }

            # Interpret security mitigations
            mitigations["ASLR"] = {
                "enabled": "DYNAMIC_BASE" in flags_found,
                "high_entropy": "HIGH_ENTROPY_VA" in flags_found,
                "description": "Address Space Layout Randomization",
            }

            mitigations["DEP"] = {
                "enabled": "NX_COMPAT" in flags_found,
                "description": "Data Execution Prevention (NX bit)",
            }

            mitigations["CFG"] = {
                "enabled": "GUARD_CF" in flags_found,
                "description": "Control Flow Guard",
            }

            mitigations["Force_Integrity"] = {
                "enabled": "FORCE_INTEGRITY" in flags_found,
                "description": "Force code integrity checks",
            }

            mitigations["AppContainer"] = {
                "enabled": "APPCONTAINER" in flags_found,
                "description": "Runs in AppContainer sandbox",
            }

            mitigations["No_SEH"] = {
                "enabled": "NO_SEH" in flags_found,
                "description": "No Structured Exception Handler",
            }

            result["mitigations"].update(mitigations)

        except Exception as e:
            logger.error(f"Error checking DllCharacteristics: {e}")

    def _check_safe_seh(self, result: dict[str, Any]):
        """Check for SafeSEH (Safe Structured Exception Handling)."""
        try:
            if self._has_no_seh_flag(result):
                self._set_safe_seh(result, False, "Binary has NO_SEH flag set")
                return

            load_config_dir = self._get_load_config_dir()
            if not load_config_dir:
                self._set_safe_seh(result, False, "No data directories found")
                return

            if self._load_config_present(load_config_dir):
                self._set_safe_seh(result, True, "Load Configuration present")
                return

            self._set_safe_seh(result, False, "No SafeSEH table found")

        except Exception as e:
            logger.error(f"Error checking SafeSEH: {e}")
            result["mitigations"]["SafeSEH"] = self._safe_seh_entry(enabled=False, error=str(e))

    def _check_load_config(self, result: dict[str, Any]):
        """Check Load Configuration Directory for security features."""
        try:
            load_config_dir = self._get_load_config_dir()
            if not load_config_dir or load_config_dir.get("vaddr", 0) == 0:
                result["load_config"] = {
                    "present": False,
                    "note": "No Load Configuration Directory",
                }
                return

            offset = load_config_dir.get("paddr", 0)
            size = load_config_dir.get("size", 0)
            if offset == 0 or size == 0:
                return

            is_64bit = self._is_pe_64bit()
            config_data = self._read_load_config_data(offset, size)
            if not config_data:
                return

            load_config = self._parse_load_config(config_data, size, offset, is_64bit, result)
            result["load_config"] = load_config

        except Exception as e:
            logger.error(f"Error checking Load Configuration: {e}")

    def _has_no_seh_flag(self, result: dict[str, Any]) -> bool:
        return result.get("dll_characteristics", {}).get("flags", {}).get("NO_SEH", False)

    def _safe_seh_entry(self, enabled: bool, note: str | None = None, error: str | None = None):
        entry = {
            "enabled": enabled,
            "description": "Safe Structured Exception Handling",
        }
        if note:
            entry["note"] = note
        if error:
            entry["error"] = error
        return entry

    def _set_safe_seh(self, result: dict[str, Any], enabled: bool, note: str):
        result["mitigations"]["SafeSEH"] = self._safe_seh_entry(enabled=enabled, note=note)

    def _get_load_config_dir(self) -> dict[str, Any] | None:
        data_dirs = silent_cmdj(self.r2, "iDj", [])
        if not data_dirs or not isinstance(data_dirs, list):
            return None
        for dd in data_dirs:
            if isinstance(dd, dict) and dd.get("name") == "LOAD_CONFIG":
                return dd
        return None

    def _load_config_present(self, load_config_dir: dict[str, Any]) -> bool:
        offset = load_config_dir.get("paddr", 0)
        if offset <= 0:
            return False
        config_data = silent_cmdj(self.r2, f"pxj 64 @ {offset}", [])
        return bool(config_data and len(config_data) >= 64)

    def _is_pe_64bit(self) -> bool:
        optional_header = silent_cmdj(self.r2, "iHj", {})
        return bool(optional_header and optional_header.get("magic", "") == "PE32+")

    def _read_load_config_data(self, offset: int, size: int) -> list[int]:
        read_size = min(size, 256)
        return silent_cmdj(self.r2, f"pxj {read_size} @ {offset}", [])

    def _parse_load_config(
        self,
        config_data: list[int],
        size: int,
        offset: int,
        is_64bit: bool,
        result: dict[str, Any],
    ) -> dict[str, Any]:
        load_config = {"present": True, "size": size, "offset": offset}
        if len(config_data) < 92:
            return load_config

        config_size = self._read_u32(config_data, 0)
        load_config["structure_size"] = config_size
        self._parse_security_cookie(load_config, config_data, is_64bit)
        self._parse_guard_flags(load_config, config_data, config_size, is_64bit, result)
        return load_config

    def _read_u32(self, data: list[int], offset: int) -> int:
        return (
            data[offset]
            | (data[offset + 1] << 8)
            | (data[offset + 2] << 16)
            | (data[offset + 3] << 24)
        )

    def _parse_security_cookie(
        self, load_config: dict[str, Any], config_data: list[int], is_64bit: bool
    ) -> None:
        cookie_offset = 88 if is_64bit else 60
        if len(config_data) <= cookie_offset + 4:
            return
        security_cookie = self._read_u32(config_data, cookie_offset)
        load_config["security_cookie"] = hex(security_cookie) if security_cookie else None

    def _parse_guard_flags(
        self,
        load_config: dict[str, Any],
        config_data: list[int],
        config_size: int,
        is_64bit: bool,
        result: dict[str, Any],
    ) -> None:
        if config_size < 148:
            return
        guard_flags_offset = 144 if is_64bit else 140
        if len(config_data) <= guard_flags_offset + 4:
            return
        guard_flags = self._read_u32(config_data, guard_flags_offset)
        load_config["guard_flags"] = {
            "value": hex(guard_flags),
            "cf_instrumented": bool(guard_flags & 0x100),
            "cf_function_table_present": bool(guard_flags & 0x200),
            "cf_long_jump_table_present": bool(guard_flags & 0x400),
            "rf_instrumented": bool(guard_flags & 0x20000),
            "rf_enable": bool(guard_flags & 0x40000),
            "rf_strict": bool(guard_flags & 0x80000),
            "retpoline_present": bool(guard_flags & 0x100000),
            "eh_continuation_table_present": bool(guard_flags & 0x400000),
        }
        self._update_mitigations_from_guard_flags(result, guard_flags)

    def _update_mitigations_from_guard_flags(
        self, result: dict[str, Any], guard_flags: int
    ) -> None:
        if guard_flags & 0x100:
            result["mitigations"]["CFG"] = {
                "enabled": True,
                "description": "Control Flow Guard",
                "details": "CF instrumented and enabled",
            }
        if guard_flags & 0x20000:
            result["mitigations"]["RFG"] = {
                "enabled": True,
                "description": "Return Flow Guard",
                "details": "RF instrumented" + (" and enabled" if guard_flags & 0x40000 else ""),
            }

    def _check_stack_cookies(self, result: dict[str, Any]):
        """Check for stack cookies (GS security)."""
        try:
            # Check if security cookie is present in Load Config
            if result.get("load_config", {}).get("security_cookie"):
                result["mitigations"]["Stack_Cookies"] = self._stack_cookie_entry(
                    enabled=True,
                    cookie=result["load_config"]["security_cookie"],
                )
                return

            # Alternative: Look for __security_cookie symbol in imports
            if self._has_security_cookie_imports():
                result["mitigations"]["Stack_Cookies"] = self._stack_cookie_entry(
                    enabled=True,
                    note="Security cookie imports found",
                )
                return

            # Check for GS-related compiler signatures
            indicator = self._find_gs_indicator()
            if indicator:
                result["mitigations"]["Stack_Cookies"] = self._stack_cookie_entry(
                    enabled=True,
                    note=f"GS indicator found: {indicator}",
                )
                return

            result["mitigations"]["Stack_Cookies"] = self._stack_cookie_entry(
                enabled=False,
                note="No stack cookie protection detected",
            )

        except Exception as e:
            logger.error(f"Error checking stack cookies: {e}")

    def _stack_cookie_entry(
        self, enabled: bool, cookie: str | None = None, note: str | None = None
    ) -> dict[str, Any]:
        entry: dict[str, Any] = {
            "enabled": enabled,
            "description": "Stack buffer overflow protection (GS)",
        }
        if cookie:
            entry["cookie"] = cookie
        if note:
            entry["note"] = note
        return entry

    def _has_security_cookie_imports(self) -> bool:
        imports = silent_cmdj(self.r2, "iij", [])
        if not imports:
            return False
        for imp in imports:
            name = imp.get("name", "").lower()
            if "__security_cookie" in name or "__security_init_cookie" in name:
                return True
        return False

    def _find_gs_indicator(self) -> str | None:
        strings = silent_cmdj(self.r2, "izzj", [])
        if not strings:
            return None
        gs_indicators = ["__GSHandlerCheck", "__security_check_cookie", "/GS"]
        for s in strings[:100]:
            string_val = s.get("string", "")
            for indicator in gs_indicators:
                if indicator in string_val:
                    return indicator
        return None

    def _check_authenticode(self, result: dict[str, Any]):
        """Check if binary is signed with Authenticode."""
        try:
            # Check for security directory
            data_dirs = silent_cmdj(self.r2, "iDj", [])
            if not data_dirs or not isinstance(data_dirs, list):
                result["mitigations"]["Authenticode"] = {
                    "enabled": False,
                    "description": "Digital signature",
                    "note": "No data directories",
                }
                return

            for dd in data_dirs:
                if (
                    isinstance(dd, dict)
                    and dd.get("name") == "SECURITY"
                    and dd.get("vaddr", 0) != 0
                ):
                    result["mitigations"]["Authenticode"] = {
                        "enabled": True,
                        "description": "Digital signature",
                        "note": "Certificate table present",
                        "size": dd.get("size", 0),
                    }
                    return

            result["mitigations"]["Authenticode"] = {
                "enabled": False,
                "description": "Digital signature",
                "note": "No certificate table found",
            }

        except Exception as e:
            logger.error(f"Error checking Authenticode: {e}")

    def _check_pe_security_features(self, result: dict[str, Any]):
        """Check additional PE header security features."""
        try:
            pe_header = self._get_pe_header()
            if not pe_header:
                return

            characteristics = pe_header.get("characteristics", 0)
            is_dll = bool(characteristics & 0x2000)
            self._check_relocations_stripped(result, characteristics)

            optional_header = silent_cmdj(self.r2, "iHj", {})
            if optional_header and isinstance(optional_header, dict):
                self._set_pe_info(result, optional_header, is_dll)
                self._check_entry_point_section(result, optional_header)

        except Exception as e:
            logger.error(f"Error checking PE security features: {e}")

    def _get_pe_header(self) -> dict[str, Any] | None:
        pe_header = silent_cmdj(self.r2, "ihj", {})
        if not pe_header:
            return None
        if isinstance(pe_header, list):
            return pe_header[0] if pe_header else None
        return pe_header if isinstance(pe_header, dict) else None

    def _check_relocations_stripped(self, result: dict[str, Any], characteristics: int) -> None:
        relocs_stripped = bool(characteristics & 0x0001)
        if relocs_stripped and result["mitigations"].get("ASLR", {}).get("enabled"):
            result["vulnerabilities"].append(
                {
                    "issue": "ASLR ineffective",
                    "description": "Relocations stripped but ASLR enabled - ASLR will not work properly",
                    "severity": "high",
                }
            )

    def _set_pe_info(self, result: dict[str, Any], optional_header: dict[str, Any], is_dll: bool):
        subsystem = optional_header.get("subsystem", "")
        if subsystem:
            result["pe_info"] = {
                "subsystem": subsystem,
                "is_dll": is_dll,
                "is_driver": "driver" in subsystem.lower(),
            }

    def _check_entry_point_section(
        self, result: dict[str, Any], optional_header: dict[str, Any]
    ) -> None:
        entry_point = optional_header.get("entry_point", 0)
        image_base = optional_header.get("image_base", 0)
        sections = silent_cmdj(self.r2, "iSj", [])
        if not (sections and isinstance(sections, list) and entry_point > 0):
            return
        entry_rva = entry_point - image_base if entry_point > image_base else entry_point
        entry_in_exec_section = self._entry_point_in_exec_section(sections, entry_rva)
        if not entry_in_exec_section:
            result["vulnerabilities"].append(
                {
                    "issue": "Unusual entry point",
                    "description": "Entry point not in executable section",
                    "severity": "medium",
                }
            )

    def _entry_point_in_exec_section(self, sections: list[dict[str, Any]], entry_rva: int) -> bool:
        for section in sections:
            if not isinstance(section, dict):
                continue
            start = section.get("vaddr", 0)
            end = start + section.get("vsize", 0)
            if start <= entry_rva < end:
                perm = section.get("perm", "")
                return perm.startswith("--x") or perm.startswith("-wx")
        return False

    def _calculate_security_score(self, result: dict[str, Any]):
        """Calculate overall security score based on enabled mitigations."""
        score = 0
        max_score = 0

        # Scoring weights for different mitigations
        mitigation_scores = {
            "ASLR": {"enabled": 15, "high_entropy": 5},
            "DEP": {"enabled": 15},
            "CFG": {"enabled": 15},
            "RFG": {"enabled": 10},
            "SafeSEH": {"enabled": 10},
            "Stack_Cookies": {"enabled": 15},
            "Authenticode": {"enabled": 10},
            "Force_Integrity": {"enabled": 5},
            "AppContainer": {"enabled": 5},
        }

        for mitigation_name, scoring in mitigation_scores.items():
            mitigation = result["mitigations"].get(mitigation_name, {})
            for check, points in scoring.items():
                max_score += points
                if mitigation.get(check):
                    score += points

        # Subtract points for vulnerabilities
        for vuln in result.get("vulnerabilities", []):
            if vuln.get("severity") == "high":
                score -= 10
            elif vuln.get("severity") == "medium":
                score -= 5

        # Ensure score doesn't go below 0
        score = max(0, score)

        result["security_score"] = {
            "score": score,
            "max_score": max_score,
            "percentage": round((score / max_score * 100) if max_score > 0 else 0, 1),
            "grade": self._get_security_grade(score, max_score),
        }

    def _get_security_grade(self, score: int, max_score: int) -> str:
        """Get security grade based on score."""
        if max_score == 0:
            return "Unknown"

        percentage = (score / max_score) * 100

        if percentage >= 90:
            return "A"
        elif percentage >= 80:
            return "B"
        elif percentage >= 70:
            return "C"
        elif percentage >= 60:
            return "D"
        else:
            return "F"

    def _generate_recommendations(self, result: dict[str, Any]):
        """Generate security recommendations based on analysis."""
        recommendations = []

        # Check each mitigation
        if not result["mitigations"].get("ASLR", {}).get("enabled"):
            recommendations.append(
                {
                    "priority": "high",
                    "mitigation": "ASLR",
                    "recommendation": "Enable ASLR by setting /DYNAMICBASE linker flag",
                    "impact": "Prevents reliable exploitation by randomizing memory addresses",
                }
            )

        if result["mitigations"].get("ASLR", {}).get("enabled") and not result["mitigations"].get(
            "ASLR", {}
        ).get("high_entropy"):
            recommendations.append(
                {
                    "priority": "medium",
                    "mitigation": "High Entropy ASLR",
                    "recommendation": "Enable high entropy ASLR with /HIGHENTROPYVA for 64-bit builds",
                    "impact": "Increases address space randomization entropy",
                }
            )

        if not result["mitigations"].get("DEP", {}).get("enabled"):
            recommendations.append(
                {
                    "priority": "high",
                    "mitigation": "DEP/NX",
                    "recommendation": "Enable DEP by setting /NXCOMPAT linker flag",
                    "impact": "Prevents code execution from data segments",
                }
            )

        if not result["mitigations"].get("CFG", {}).get("enabled"):
            recommendations.append(
                {
                    "priority": "high",
                    "mitigation": "Control Flow Guard",
                    "recommendation": "Enable CFG with /guard:cf compiler and linker flags",
                    "impact": "Prevents control flow hijacking attacks",
                }
            )

        if not result["mitigations"].get("Stack_Cookies", {}).get("enabled"):
            recommendations.append(
                {
                    "priority": "high",
                    "mitigation": "Stack Cookies",
                    "recommendation": "Enable stack cookies with /GS compiler flag",
                    "impact": "Detects stack buffer overflows",
                }
            )

        if not result["mitigations"].get("SafeSEH", {}).get("enabled"):
            # Only recommend for 32-bit binaries
            if not result.get("pe_info", {}).get("is_64bit", True):
                recommendations.append(
                    {
                        "priority": "medium",
                        "mitigation": "SafeSEH",
                        "recommendation": "Enable SafeSEH with /SAFESEH linker flag (32-bit only)",
                        "impact": "Prevents SEH overwrites",
                    }
                )

        if not result["mitigations"].get("Authenticode", {}).get("enabled"):
            recommendations.append(
                {
                    "priority": "medium",
                    "mitigation": "Code Signing",
                    "recommendation": "Sign the binary with Authenticode certificate",
                    "impact": "Ensures binary integrity and authenticity",
                }
            )

        result["recommendations"] = recommendations
