http_proxy.py

#!/usr/bin/env python3
"""
HTTP/HTTPS proxy with SNI inspection for AlwaysBlock
Handles both HTTP and HTTPS CONNECT requests
"""
import socket
import select
import threading
import logging
import json
import resource
import time
import urllib.request
from pathlib import Path

logger = logging.getLogger(__name__)


def raise_fd_limit(target=65536):
    """Raise the soft RLIMIT_NOFILE so bursts of concurrent connections don't
    exhaust file descriptors. macOS launchd starts daemons with a soft limit
    of 256; each proxied connection uses two sockets, so a single busy page
    can blow past that and cause both EMFILE on accept() and EAI_NONAME on
    getaddrinfo() (since the resolver can't allocate a query socket either).
    """
    try:
        soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
        new_soft = min(target, hard) if hard != resource.RLIM_INFINITY else target
        if new_soft > soft:
            resource.setrlimit(resource.RLIMIT_NOFILE, (new_soft, hard))
            logger.info(f"Raised RLIMIT_NOFILE soft limit: {soft} -> {new_soft} (hard={hard})")
        else:
            logger.info(f"RLIMIT_NOFILE already adequate: soft={soft}, hard={hard}")
    except (ValueError, OSError) as e:
        logger.warning(f"Could not raise RLIMIT_NOFILE: {e}")


class HTTPProxy:
    """HTTP/HTTPS proxy that blocks based on hostname"""

    def __init__(self, port=8905, blocked_domains_file="/tmp/alwaysblock_domains.json"):
        self.port = port
        self.blocked_domains_file = Path(blocked_domains_file)
        self.blocked_domains = set()
        self.excluded_domains = set()
        self.server_socket = None
        self.running = False
        self.last_mtime = 0
        self.captive_portal_mode = False
        self.captive_portal_entered_at = 0
        self.last_captive_check = 0
        self.pause_until = 0  # Manual pause expiry timestamp
        self.recent_failures = 0  # Consecutive upstream failures (reset on success)
        self.first_recent_failure_at = 0

    def load_blocked_domains(self):
        """Load blocked domains from JSON file"""
        try:
            if self.blocked_domains_file.exists():
                with open(self.blocked_domains_file, 'r') as f:
                    data = json.load(f)
                    self.blocked_domains = set(data.get('domains', []))
                    self.excluded_domains = set(data.get('excluded', []))
                    self.pause_until = data.get('pause_until', 0)
                    self.last_mtime = self.blocked_domains_file.stat().st_mtime
                    logger.info(f"Loaded {len(self.blocked_domains)} blocked domains, {len(self.excluded_domains)} excluded")
            else:
                logger.warning(f"Blocked domains file not found: {self.blocked_domains_file}")
                self.blocked_domains = set()
                self.excluded_domains = set()
        except Exception as e:
            logger.error(f"Failed to load blocked domains: {e}")
            self.blocked_domains = set()
            self.excluded_domains = set()

    def check_and_reload(self):
        """Check if domains file changed and reload if needed"""
        try:
            if self.blocked_domains_file.exists():
                mtime = self.blocked_domains_file.stat().st_mtime
                if mtime > self.last_mtime:
                    self.load_blocked_domains()
        except Exception as e:
            logger.debug(f"Error checking file mtime: {e}")

    def check_captive_portal(self):
        """Check if we're behind a captive portal"""
        try:
            # Bypass proxy to avoid loop (we are the proxy)
            proxy_handler = urllib.request.ProxyHandler({})
            opener = urllib.request.build_opener(proxy_handler)
            req = opener.open('http://captive.apple.com/hotspot-detect.html', timeout=5)
            is_captive = b'Success' not in req.read()
            if is_captive != self.captive_portal_mode:
                self.captive_portal_mode = is_captive
                if is_captive:
                    self.captive_portal_entered_at = time.time()
                    logger.info("📶 Captive portal detected - allowing all traffic for 2 min")
                else:
                    self.captive_portal_entered_at = 0
                    logger.info("✅ Internet connected - resuming blocking")
        except Exception:
            if not self.captive_portal_mode:
                self.captive_portal_mode = True
                self.captive_portal_entered_at = time.time()
                logger.info("📶 Network issue detected - allowing all traffic for 2 min")

    def record_connection_failure(self):
        """Record a connection failure and check for captive portal.

        A single transient DNS/connect blip should NOT disable blocking. We only
        run the captive-portal check after several failures cluster together
        (a real network outage), so one bad lookup can't turn off blocking for 2 min.
        """
        now = time.time()
        # Reset the cluster if the last failure was a while ago
        if now - self.first_recent_failure_at > 15:
            self.recent_failures = 0
            self.first_recent_failure_at = now
        self.recent_failures += 1

        # Require a cluster of failures before suspecting a real network problem
        if self.recent_failures < 3:
            return

        # Rate-limited to once per 10s
        if now - self.last_captive_check > 10:
            self.last_captive_check = now
            self.check_captive_portal()

    def record_connection_success(self):
        """Clear the failure cluster after a successful upstream connection."""
        self.recent_failures = 0

    def should_block_domain(self, hostname):
        """Check if domain should be blocked (with subdomain matching)"""
        if not hostname:
            return False

        now = time.time()

        # Check manual pause (from pause command)
        if self.pause_until > now:
            return False

        # Check captive portal mode with 2-minute auto-expire
        if self.captive_portal_mode:
            if self.captive_portal_entered_at > 0 and (now - self.captive_portal_entered_at) > 120:
                # Auto-expire after 2 minutes
                self.captive_portal_mode = False
                self.captive_portal_entered_at = 0
                logger.info("✅ Captive portal mode expired - resuming blocking")
            else:
                return False

        # First check if domain is explicitly excluded (takes precedence)
        if hostname in self.excluded_domains:
            return False

        # Check without www prefix for exclusions
        if hostname.startswith('www.'):
            if hostname[4:] in self.excluded_domains:
                return False

        # Check if any parent domain is excluded
        parts = hostname.split('.')
        for i in range(len(parts)):
            domain = '.'.join(parts[i:])
            if domain in self.excluded_domains:
                return False

        # Now check if domain should be blocked
        # Direct match
        if hostname in self.blocked_domains:
            return True

        # Check without www prefix
        if hostname.startswith('www.'):
            if hostname[4:] in self.blocked_domains:
                return True

        # Check parent domains (subdomain matching)
        for i in range(len(parts)):
            domain = '.'.join(parts[i:])
            if domain in self.blocked_domains:
                return True

        return False

    def connect_upstream(self, hostname, port):
        """Resolve and connect to the upstream server.

        Resilient to transient DNS blips: getaddrinfo is retried a couple of
        times, and every resolved address is tried (not just the first), so a
        flaky resolver or one dead IP doesn't turn into an instant 502.
        Returns a connected socket, or None on failure.
        """
        addrinfo = None
        last_err = None
        for attempt in range(3):
            try:
                addrinfo = socket.getaddrinfo(hostname, port, socket.AF_INET, socket.SOCK_STREAM)
                break
            except socket.gaierror as e:
                last_err = e
                # Transient resolver failure (EAI_NONAME/EAI_AGAIN etc.). Brief backoff, then retry.
                time.sleep(0.15 * (attempt + 1))
        if not addrinfo:
            logger.error(f"DNS resolution failed for {hostname}:{port}: {last_err}")
            return None

        for family, socktype, proto, _canon, sockaddr in addrinfo:
            server_socket = socket.socket(family, socktype, proto)
            server_socket.settimeout(5.0)
            try:
                server_socket.connect(sockaddr)
                self.record_connection_success()
                return server_socket
            except Exception as e:
                last_err = e
                try:
                    server_socket.close()
                except Exception:
                    pass
        logger.error(f"Failed to connect to {hostname}:{port}: {last_err}")
        return None

    def handle_connection(self, client_socket, client_address):
        """Handle a single client connection"""
        try:
            client_socket.settimeout(5.0)

            # Read the HTTP request (read up to 8KB or until we have the first line)
            buffer = client_socket.recv(8192)
            if not buffer:
                client_socket.close()
                return

            # Split to get first line
            lines = buffer.split(b'\r\n')
            request_line = lines[0].decode('utf-8', errors='ignore').strip()

            # Parse the request
            parts = request_line.split(' ')
            if len(parts) < 2:
                client_socket.close()
                return

            method = parts[0]
            url = parts[1]

            # Extract hostname
            if method == 'CONNECT':
                # HTTPS proxy request: CONNECT example.com:443 HTTP/1.1
                hostname = url.split(':')[0]
                port = int(url.split(':')[1]) if ':' in url else 443
            else:
                # HTTP proxy request: GET http://example.com/path HTTP/1.1
                if url.startswith('http://'):
                    url_parts = url[7:].split('/', 1)
                    hostname = url_parts[0].split(':')[0]
                    port = int(url_parts[0].split(':')[1]) if ':' in url_parts[0] else 80
                else:
                    # Relative URL, need to find Host header in the buffer we already read
                    hostname = None
                    for line in lines[1:]:
                        line_str = line.decode('utf-8', errors='ignore').strip()
                        if line_str.lower().startswith('host:'):
                            hostname = line_str.split(':', 1)[1].strip()
                            break

                    if not hostname:
                        client_socket.close()
                        return
                    port = 80

            # Check if blocked
            if self.should_block_domain(hostname):
                logger.info(f"🚫 Blocking {method} to: {hostname}")

                # Send error response
                if method == 'CONNECT':
                    response = b'HTTP/1.1 403 Forbidden\r\n\r\n'
                else:
                    response = b'HTTP/1.1 403 Forbidden\r\nContent-Type: text/html\r\n\r\n<html><body><h1>Blocked by AlwaysBlock</h1></body></html>'

                try:
                    client_socket.sendall(response)
                except:
                    pass

                client_socket.close()
                return

            # Allow the connection
            logger.debug(f"✅ Allowing {method} to: {hostname}:{port}")

            # Connect to the real server (with DNS retry + multi-IP fallback)
            server_socket = self.connect_upstream(hostname, port)
            if server_socket is None:
                self.record_connection_failure()

                response = b'HTTP/1.1 502 Bad Gateway\r\n\r\n'
                try:
                    client_socket.sendall(response)
                except:
                    pass

                client_socket.close()
                return

            if method == 'CONNECT':
                # HTTPS: Send 200 OK then just forward bytes
                client_socket.sendall(b'HTTP/1.1 200 Connection Established\r\n\r\n')

                # Remove timeouts
                client_socket.settimeout(None)
                server_socket.settimeout(None)

                # Forward data bidirectionally
                self.forward_data(client_socket, server_socket)
            else:
                # HTTP: Forward the request and handle response
                # Send the original request to the server
                server_socket.sendall(request_line.encode() + b'\r\n')

                # Forward remaining headers and body
                while True:
                    data = client_socket.recv(8192)
                    if not data:
                        break
                    server_socket.sendall(data)
                    if len(data) < 8192:
                        break

                # Forward response back
                while True:
                    data = server_socket.recv(8192)
                    if not data:
                        break
                    client_socket.sendall(data)

                client_socket.close()
                server_socket.close()

        except socket.timeout:
            logger.debug(f"Connection timeout from {client_address[0]}")
        except Exception as e:
            logger.error(f"Error handling connection: {e}")
        finally:
            try:
                client_socket.close()
            except:
                pass

    def forward_data(self, client_socket, server_socket):
        """Forward data bidirectionally between sockets"""
        try:
            sockets = [client_socket, server_socket]

            while True:
                readable, _, exceptional = select.select(sockets, [], sockets, 60.0)

                if exceptional:
                    break

                if not readable:
                    break

                for sock in readable:
                    try:
                        data = sock.recv(8192)
                        if not data:
                            return

                        if sock is client_socket:
                            server_socket.sendall(data)
                        else:
                            client_socket.sendall(data)

                    except Exception as e:
                        logger.debug(f"Error forwarding data: {e}")
                        return

        except Exception as e:
            logger.debug(f"Connection forwarding error: {e}")
        finally:
            try:
                client_socket.close()
            except:
                pass
            try:
                server_socket.close()
            except:
                pass

    def start(self):
        """Start the HTTP proxy server"""
        try:
            self.load_blocked_domains()
            self.check_captive_portal()

            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.server_socket.bind(('127.0.0.1', self.port))
            self.server_socket.listen(128)

            self.running = True
            logger.info(f"🚀 HTTP proxy started on 127.0.0.1:{self.port}")
            logger.info(f"📋 Blocking {len(self.blocked_domains)} domains")

            # Start background thread to check for file changes and captive portal
            def reload_checker():
                while self.running:
                    time.sleep(5)
                    self.check_and_reload()
                    # Check captive portal: every 5s if in captive mode (to resume blocking fast)
                    # Otherwise rely on connection failure detection, with 60s fallback
                    now = time.time()
                    interval = 5 if self.captive_portal_mode else 60
                    if now - self.last_captive_check >= interval:
                        self.check_captive_portal()
                        self.last_captive_check = now

            checker_thread = threading.Thread(target=reload_checker, daemon=True)
            checker_thread.start()

            while self.running:
                try:
                    self.server_socket.settimeout(1.0)
                    try:
                        client_socket, client_address = self.server_socket.accept()
                    except socket.timeout:
                        continue

                    thread = threading.Thread(
                        target=self.handle_connection,
                        args=(client_socket, client_address),
                        daemon=True
                    )
                    thread.start()

                except KeyboardInterrupt:
                    logger.info("Shutting down proxy...")
                    break
                except Exception as e:
                    logger.error(f"Error accepting connection: {e}")

        except Exception as e:
            logger.error(f"Failed to start proxy: {e}")
        finally:
            self.stop()

    def stop(self):
        """Stop the proxy server"""
        self.running = False
        if self.server_socket:
            try:
                self.server_socket.close()
            except:
                pass
        logger.info("Proxy stopped")


if __name__ == '__main__':
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s [%(levelname)s] %(message)s'
    )

    raise_fd_limit()

    proxy = HTTPProxy()
    proxy.start()