"""Enhanced SSRF attack protection with DNS pinning and rebinding prevention."""

# Standard library imports
import ipaddress
import socket
import time
from typing import Any, Dict, List, Optional, Set, Tuple

# Third-party imports
import structlog

logger = structlog.get_logger(__name__)


class SSRFProtection:
    """Enhanced SSRF attack protection with DNS pinning."""

    def __init__(self) -> None:
        """Initialize SSRF protection with blocked IPs and networks."""
        self._dns_cache: Dict[str, Tuple[str, float, int]] = {}
        self._dns_cache_ttl = 300

        self._blocked_ips: Set[str] = {
            "127.0.0.1",
            "::1",
            "0.0.0.0",
            "169.254.169.254",
            "10.0.0.1",
            "10.255.255.255",
            "172.16.0.1",
            "172.31.255.255",
            "192.168.0.1",
            "192.168.255.255",
        }

        self._blocked_networks = [
            "127.0.0.0/8",
            "169.254.0.0/16",
            "224.0.0.0/4",
            "240.0.0.0/4",
            "172.17.0.0/16",
            "172.18.0.0/16",
            "172.19.0.0/16",
            "172.20.0.0/16",
            "172.21.0.0/16",
            "172.22.0.0/16",
            "172.23.0.0/16",
            "172.24.0.0/16",
            "172.25.0.0/16",
            "172.26.0.0/16",
            "172.27.0.0/16",
            "172.28.0.0/16",
            "172.29.0.0/16",
            "172.30.0.0/16",
            "172.31.0.0/16",
        ]

        self._allowed_networks: Dict[str, List[str]] = {
            "USER": ["10.0.0.0/8", "172.16.0.0/16", "192.168.1.0/24"],
            "ADMIN": [
                "10.0.0.0/8",
                "172.16.0.0/16",
                "192.168.0.0/16",
                "203.0.113.0/24",
            ],
            "SUPER_ADMIN": ["0.0.0.0/0"],
        }
    
    def validate_host(
        self, hostname: str, user_role: str
    ) -> Tuple[bool, str]:
        """Validate host with enhanced SSRF protection.

        Args:
            hostname: Hostname or IP address
            user_role: User role

        Returns:
            Tuple of (allowed: bool, reason: str)
        """
        try:
            if not hostname or len(hostname) > 253:
                return False, f"Invalid hostname length: {hostname}"

            suspicious_chars = [
                "..",
                "//",
                "\\",
                "<",
                ">",
                '"',
                "'",
                "`",
                "\x00",
            ]
            if any(char in hostname for char in suspicious_chars):
                return False, f"Suspicious characters in hostname: {hostname}"

            if hostname.lower() in ("localhost", "127.0.0.1", "::1"):
                return False, f"Host {hostname} is blocked (localhost)"

            resolved_ip = self._resolve_hostname_with_pinning(hostname)
            if not resolved_ip:
                return False, f"Cannot resolve hostname: {hostname}"

            if resolved_ip in self._blocked_ips:
                return False, f"IP {resolved_ip} is in blocked list"

            ip_addr = ipaddress.ip_address(resolved_ip)
            for blocked_network in self._blocked_networks:
                if ip_addr in ipaddress.ip_network(blocked_network):
                    return (
                        False,
                        f"IP {resolved_ip} is in blocked network {blocked_network}",
                    )

            allowed_networks = self._allowed_networks.get(user_role, [])
            if not allowed_networks:
                return False, f"Role {user_role} has no allowed networks"

            if user_role == "SUPER_ADMIN":
                return True, f"IP {resolved_ip} allowed for SUPER_ADMIN"

            for allowed_network in allowed_networks:
                if ip_addr in ipaddress.ip_network(allowed_network):
                    return (
                        True,
                        f"IP {resolved_ip} allowed in network {allowed_network}",
                    )

            return (
                False,
                f"IP {resolved_ip} not in any allowed network for role {user_role}",
            )

        except Exception as e:
            logger.error("SSRF validation error", hostname=hostname, error=str(e))
            return False, f"Error validating host: {str(e)}"
    
    def _resolve_hostname_with_pinning(self, hostname: str) -> Optional[str]:
        """DNS resolution with pinning to prevent rebinding attacks.

        Args:
            hostname: Hostname to resolve

        Returns:
            IP address or None if resolution failed
        """
        try:
            cache_key = hostname.lower()
            if cache_key in self._dns_cache:
                cached_ip, timestamp, ttl = self._dns_cache[cache_key]

                if time.time() - timestamp < ttl:
                    logger.debug(
                        "Using cached DNS resolution",
                        hostname=hostname,
                        ip=cached_ip,
                        age_seconds=int(time.time() - timestamp),
                    )
                    return cached_ip
                del self._dns_cache[cache_key]

            original_timeout = socket.getdefaulttimeout()
            socket.setdefaulttimeout(5)

            try:
                ip1 = socket.gethostbyname(hostname)

                time.sleep(0.1)

                ip2 = socket.gethostbyname(hostname)

                if ip1 != ip2:
                    logger.warning(
                        "DNS rebinding detected",
                        hostname=hostname,
                        ip1=ip1,
                        ip2=ip2,
                    )
                    return None

                if ip1 in ("127.0.0.1", "::1"):
                    logger.warning(
                        "DNS resolution returned localhost", hostname=hostname, ip=ip1
                    )
                    return None

                self._dns_cache[cache_key] = (
                    ip1,
                    time.time(),
                    self._dns_cache_ttl,
                )

                logger.info(
                    "DNS resolution successful", hostname=hostname, ip=ip1, cached=True
                )

                return ip1

            finally:
                socket.setdefaulttimeout(original_timeout)

        except socket.gaierror as e:
            logger.warning("DNS resolution failed", hostname=hostname, error=str(e))
            return None
        except socket.timeout:
            logger.warning("DNS resolution timeout", hostname=hostname)
            return None
        except Exception as e:
            logger.error(
                "Unexpected DNS resolution error", hostname=hostname, error=str(e)
            )
            return None
    
    def validate_port(self, port: int) -> Tuple[bool, str]:
        """Validate port number.

        Args:
            port: Port number

        Returns:
            Tuple of (valid: bool, reason: str)
        """
        if not isinstance(port, int) or port < 1 or port > 65535:
            return False, f"Invalid port: {port}"

        blocked_ports = {
            22,
            23,
            25,
            53,
            80,
            110,
            143,
            443,
            993,
            995,
            135,
            139,
            445,
            1433,
            1521,
            3306,
            5432,
            6379,
            3389,
            5900,
            5901,
            5902,
            5903,
            5904,
            5905,
            8080,
            8443,
            9090,
            9091,
        }

        if port in blocked_ports:
            return False, f"Port {port} is blocked (system port)"

        return True, f"Port {port} is valid"
    
    def cleanup_expired_cache(self) -> None:
        """Clean up expired DNS cache entries."""
        current_time = time.time()
        expired_keys = [
            key
            for key, (_, timestamp, ttl) in self._dns_cache.items()
            if current_time - timestamp > ttl
        ]

        for key in expired_keys:
            del self._dns_cache[key]

        if expired_keys:
            logger.info(
                "Cleaned up expired DNS cache entries", count=len(expired_keys)
            )

    def get_cache_stats(self) -> Dict[str, Any]:
        """Get DNS cache statistics.

        Returns:
            Dictionary with cache statistics
        """
        current_time = time.time()
        active_entries = 0
        expired_entries = 0

        for _, timestamp, ttl in self._dns_cache.values():
            if current_time - timestamp < ttl:
                active_entries += 1
            else:
                expired_entries += 1

        return {
            "total_entries": len(self._dns_cache),
            "active_entries": active_entries,
            "expired_entries": expired_entries,
            "cache_ttl_seconds": self._dns_cache_ttl,
            "blocked_ips_count": len(self._blocked_ips),
            "blocked_networks_count": len(self._blocked_networks),
        }

    def add_blocked_ip(self, ip: str) -> None:
        """Add IP to blocked list.

        Args:
            ip: IP address to block
        """
        self._blocked_ips.add(ip)
        logger.info("Added IP to blocked list", ip=ip)

    def remove_blocked_ip(self, ip: str) -> None:
        """Remove IP from blocked list.

        Args:
            ip: IP address to unblock
        """
        self._blocked_ips.discard(ip)
        logger.info("Removed IP from blocked list", ip=ip)


# Global instance for use in API
ssrf_protection = SSRFProtection()
