← Back
Editing: net_transport.py
"""Networking transport helpers for urllib. This module provides a small abstraction on top of urllib.request so that callers can keep using urllib.request.Request, but routing of connections can be customized: - hostname resolution is handled in user code; - selected IP may be randomized or chosen using any complex logic; - for HTTPS: connects to a chosen IP but keeps correct SNI and certificate hostname validation for the original hostname (NOT the IP). Examples: Default behavior (plain urllib): from defence360agent.utils.net_transport import UrlTransport transport = UrlTransport() req = urllib.request.Request( "https://files.imunify360.com/static/sigs/v1/description.json" ) with transport.open(req, timeout=10) as resp: body = resp.read() Randomize target IP on each connection (A/AAAA -> random choice): from defence360agent.utils.net_transport import UrlTransport, RandomIpChooser chooser = RandomIpChooser() transport = UrlTransport(ip_chooser=chooser) req = urllib.request.Request( "https://files.imunify360.com/static/sigs/v1/description.json" ) with transport.open(req, timeout=10) as resp: body = resp.read() Notes: - HTTPS: connects to the chosen IP but keeps SNI/cert checks against original hostname. - HTTP: Host header stays original hostname because urllib builds it from the URL. """ import http.client import ipaddress import random import socket import threading import time import urllib.request from abc import ABC, abstractmethod from logging import getLogger from typing import Dict, Optional, Tuple, TYPE_CHECKING if TYPE_CHECKING: import ssl logger = getLogger(__name__) #: default cache TTL for DNS responses _DNS_DEFAULT_TTL_SECONDS = 300.0 def _is_ipv4(ip: str) -> bool: """Return True if *ip* is an IPv4 address string. Implementation relies solely on ipaddress.ip_address for correctness. """ try: return isinstance(ipaddress.ip_address(ip), ipaddress.IPv4Address) except ValueError: return False class IpChooser(ABC): """Select an IP address to connect to for a given hostname and port. Implementations may be stateful and can keep caches/metrics inside. """ @abstractmethod def choose(self, hostname: str, port: int) -> str: """Return an IP address (v4 or v6) for *hostname*:*port*.""" raise NotImplementedError def __call__(self, hostname: str, port: int) -> str: return self.choose(hostname, port) class DnsCacheResolver: """DNS cache for socket.getaddrinfo() results. It caches per (hostname, port, family). This is intentionally small and local: it is meant only to avoid excessive getaddrinfo() calls. """ def __init__( self, *, family: int = socket.AF_UNSPEC, ttl_seconds: float = _DNS_DEFAULT_TTL_SECONDS, ): self._family = family self._ttl_seconds = ttl_seconds self._cache: Dict[ Tuple[str, int, int], Tuple[float, Tuple[str, ...]] ] = {} self._lock = threading.Lock() def get_ips(self, hostname: str, port: int) -> Tuple[str, ...]: key = (hostname, port, self._family) now = time.time() with self._lock: cached = self._cache.get(key) if cached is not None: expires_at, ips = cached if now < expires_at: logger.debug( "DnsCacheResolver cache hit for %s:%s (family=%s)", hostname, port, self._family, ) return ips logger.debug( "DnsCacheResolver cache miss/expired for %s:%s (family=%s)", hostname, port, self._family, ) infos = socket.getaddrinfo( hostname, port, self._family, socket.SOCK_STREAM, ) ips = [] for _, _, _, _, sockaddr in infos: ip = sockaddr[0] if ip not in ips: ips.append(ip) if not ips: raise OSError("No IPs resolved for {}:{}".format(hostname, port)) ips_t = tuple(ips) with self._lock: self._cache[key] = (now + self._ttl_seconds, ips_t) logger.debug( "DnsCacheResolver resolved %s:%s (family=%s) to %s", hostname, port, self._family, ips_t, ) return ips_t class RandomIpChooserWithIPv6Toggle(IpChooser): """Resolve hostname and select a random IP. IPv6 selection can be enabled/disabled at runtime: - when IPv6 is enabled: choose from IPv4 + IPv6 candidates - when IPv6 is disabled: choose from IPv4-only candidates """ def __init__( self, *, resolver: Optional[DnsCacheResolver] = None, rng: Optional[random.Random] = None, ipv6_enabled: bool = True, ): self._resolver = resolver or DnsCacheResolver() self._rng = rng or random.Random() self._ipv6_enabled = ipv6_enabled self._last_ip: Optional[str] = None def enable_ipv6(self) -> None: self._ipv6_enabled = True def disable_ipv6(self) -> None: self._ipv6_enabled = False def is_ipv6_enabled(self) -> bool: return self._ipv6_enabled def last_ip(self) -> Optional[str]: return self._last_ip def last_ip_was_ipv6(self) -> bool: return bool(self._last_ip) and (":" in self._last_ip) def choose(self, hostname: str, port: int) -> str: ips = self._resolver.get_ips(hostname, port) if self._ipv6_enabled: chosen = self._rng.choice(ips) self._last_ip = chosen logger.debug( "RandomIpChooserWithIPv6Toggle selected IP %s for %s:%s " "(IPv6 enabled)", chosen, hostname, port, ) return chosen ipv4_ips = tuple(ip for ip in ips if _is_ipv4(ip)) if not ipv4_ips: raise OSError( "No IPv4 IPs resolved for {}:{}".format(hostname, port) ) chosen = self._rng.choice(ipv4_ips) self._last_ip = chosen logger.debug( "RandomIpChooserWithIPv6Toggle selected IPv4 IP %s for %s:%s " "(IPv6 disabled)", chosen, hostname, port, ) return chosen class RandomIpChooser(IpChooser): """Resolve hostname and select a random IP from resolved candidates.""" def __init__( self, *, resolver: Optional[DnsCacheResolver] = None, rng: Optional[random.Random] = None, ): self._resolver = resolver or DnsCacheResolver() self._rng = rng or random.Random() def choose(self, hostname: str, port: int) -> str: ips = self._resolver.get_ips(hostname, port) chosen = self._rng.choice(ips) logger.debug( "RandomIpChooser selected IP %s for %s:%s", chosen, hostname, port, ) return chosen class ForcedIPHTTPConnection(http.client.HTTPConnection): """HTTPConnection that connects to a chosen IP. Important: urllib builds the request URL with the original hostname, therefore the Host header stays correct. """ def __init__( self, hostname: str, port: Optional[int] = None, *, ip_chooser: IpChooser, timeout=socket._GLOBAL_DEFAULT_TIMEOUT, source_address=None, ): super().__init__( hostname, port=port, timeout=timeout, source_address=source_address, ) self._ip_chooser = ip_chooser def connect(self) -> None: port = self.port or 80 ip = self._ip_chooser.choose(self.host, port) logger.debug( "ForcedIPHTTPConnection connecting to %s:%s for hostname %s", ip, port, self.host, ) self.sock = socket.create_connection( (ip, port), self.timeout, self.source_address, ) class ForcedIPHTTPSConnection(http.client.HTTPSConnection): """HTTPSConnection that connects to a chosen IP. TLS details: - Uses original hostname for SNI (server_hostname in wrap_socket) - Certificate hostname validation is performed for the original hostname """ def __init__( self, hostname: str, port: Optional[int] = None, *, ip_chooser: IpChooser, context: "ssl.SSLContext", timeout=socket._GLOBAL_DEFAULT_TIMEOUT, source_address=None, ): super().__init__( hostname, port=port, context=context, timeout=timeout, source_address=source_address, ) self._ip_chooser = ip_chooser def connect(self) -> None: port = self.port or 443 ip = self._ip_chooser.choose(self.host, port) logger.debug( "ForcedIPHTTPSConnection connecting to %s:%s for hostname %s", ip, port, self.host, ) raw_sock = socket.create_connection( (ip, port), self.timeout, self.source_address, ) if self._tunnel_host: self.sock = raw_sock self._tunnel() raw_sock = self.sock self.sock = self._context.wrap_socket( raw_sock, server_hostname=self.host, ) class ForcedIPHTTPHandler(urllib.request.HTTPHandler): """urllib handler that creates ForcedIPHTTPConnection.""" def __init__(self, *, ip_chooser: IpChooser): super().__init__() self._ip_chooser = ip_chooser def http_open(self, req) -> http.client.HTTPResponse: def factory(host, **kwargs): return ForcedIPHTTPConnection( host, ip_chooser=self._ip_chooser, timeout=kwargs.get("timeout"), ) return self.do_open(factory, req) class ForcedIPHTTPSHandler(urllib.request.HTTPSHandler): """urllib handler that creates ForcedIPHTTPSConnection.""" def __init__(self, *, ip_chooser: IpChooser, context: "ssl.SSLContext"): super().__init__(context=context) self._ip_chooser = ip_chooser self._context = context def https_open(self, req) -> http.client.HTTPResponse: def factory(host, **kwargs): return ForcedIPHTTPSConnection( host, ip_chooser=self._ip_chooser, context=self._context, timeout=kwargs.get("timeout"), ) return self.do_open(factory, req) class UrlTransport: """Single entrypoint for opening urllib requests. If *ip_chooser* is provided, the transport will connect to the selected IP address, while keeping correct Host/SNI/cert validation for the original hostname. If *ip_chooser* is not provided, it behaves like plain urllib. """ def __init__( self, *, ip_chooser: Optional[IpChooser] = None, ssl_context: Optional["ssl.SSLContext"] = None, ): import ssl as _ssl self._ssl_context = ssl_context or _ssl.create_default_context() if ip_chooser is None: self._opener = urllib.request.build_opener() else: self._opener = urllib.request.build_opener( ForcedIPHTTPHandler(ip_chooser=ip_chooser), ForcedIPHTTPSHandler( ip_chooser=ip_chooser, context=self._ssl_context, ), ) def open( self, req: urllib.request.Request, *, timeout: Optional[float] = None, ) -> http.client.HTTPResponse: if timeout is None: return self._opener.open(req) return self._opener.open(req, timeout=timeout)
Save File
Cancel