"""Strict outbound payload fetcher. HTTP(S) GET only; never executes, pipes or writes files.

Egress policy: only http/https on ports 80/443; the host must resolve exclusively to global
unicast addresses (private/loopback/link-local/metadata rejected); the validated IP is pinned
for the actual connect so DNS rebinding cannot redirect it; redirects are re-validated and
depth-limited; size and time are bounded. Downloaded bytes go only to inert evidence storage.
"""
import http.client
import ipaddress
import re
import socket
import ssl
import threading
import urllib.parse

URL_RE = re.compile(r'https?://[^\s\'"`;|&$()<>\\]+', re.I)
MAX_BYTES = 16 * 1024 * 1024
TIMEOUT = 10
ALLOWED_PORTS = (80, 443)


def extract_urls(text):
    return URL_RE.findall(text or '')


def public_addresses(host):
    infos = socket.getaddrinfo(host, None, proto=socket.IPPROTO_TCP)
    addresses = []
    for info in infos:
        address = info[4][0]
        ip = ipaddress.ip_address(address)
        if not ip.is_global or ip.is_multicast:
            raise ValueError('host resolves to a non-public address')
        if address not in addresses:
            addresses.append(address)
    if not addresses:
        raise ValueError('host did not resolve')
    return addresses


class _PinnedHTTP(http.client.HTTPConnection):
    def __init__(self, host, ip, port):
        super().__init__(host, port, timeout=TIMEOUT)
        self._ip = ip

    def connect(self):
        self.sock = socket.create_connection((self._ip, self.port), self.timeout)


class _PinnedHTTPS(http.client.HTTPSConnection):
    def __init__(self, host, ip, port):
        super().__init__(host, port, timeout=TIMEOUT, context=ssl.create_default_context())
        self._ip = ip

    def connect(self):
        raw = socket.create_connection((self._ip, self.port), self.timeout)
        self.sock = self._context.wrap_socket(raw, server_hostname=self.host)


class Fetcher:
    def __init__(self, emit, resolver=public_addresses, max_bytes=MAX_BYTES,
                 ports=ALLOWED_PORTS, redirects=3, slots=4, on_failure=None):
        self.emit, self.resolver, self.on_failure = emit, resolver, on_failure
        self.max_bytes, self.ports, self.redirects = max_bytes, ports, redirects
        self.slots = threading.BoundedSemaphore(slots)

    def submit(self, text, context):
        for url in list(dict.fromkeys(extract_urls(text)))[:4]:
            if not self.slots.acquire(False):
                self._safe_emit('payload_fetch_rejected', url=url, reason='concurrency', **context)
                continue
            threading.Thread(target=self._run, args=(url, context), daemon=True).start()

    def _safe_emit(self, event, **data):
        try:
            self.emit(event, **data)
        except OSError as error:
            if self.on_failure:
                self.on_failure(error)

    def _run(self, url, context):
        try:
            self._fetch(url, context)
        except Exception as error:
            self._safe_emit('payload_fetch_error', url=url, error=type(error).__name__, **context)
        finally:
            self.slots.release()

    def _fetch(self, url, context, depth=0):
        parts = urllib.parse.urlsplit(url)
        if parts.scheme not in ('http', 'https') or not parts.hostname:
            raise ValueError('blocked scheme or host')
        port = parts.port or (443 if parts.scheme == 'https' else 80)
        if port not in self.ports:
            raise ValueError('blocked port')
        addresses = self.resolver(parts.hostname)
        connection = (_PinnedHTTPS if parts.scheme == 'https' else _PinnedHTTP)(
            parts.hostname, addresses[0], port)
        path = urllib.parse.urlunsplit(('', '', parts.path or '/', parts.query, ''))
        try:
            connection.request('GET', path, headers={'Host': parts.hostname,
                               'User-Agent': 'CitrixGateway', 'Accept': '*/*', 'Connection': 'close'})
            response = connection.getresponse()
            data = response.read(self.max_bytes + 1)
            status, location = response.status, response.getheader('Location')
            content_type = (response.getheader('Content-Type') or '')[:200]
        finally:
            connection.close()
        self._safe_emit('payload_download', url=url, status=status, content_type=content_type,
                        truncated=len(data) > self.max_bytes, content=data[:self.max_bytes], **context)
        # ponytail: only 3 redirect hops; raise the cap if a real payload chain needs more.
        if location and status in (301, 302, 303, 307, 308) and depth < self.redirects:
            self._fetch(urllib.parse.urljoin(url, location), context, depth + 1)