proxy.py

#!/usr/bin/env python3
"""
Caching HTTP Proxy Server

A non-blocking, select()-based HTTP proxy with disk-based response caching.

URL format:
    http://localhost:<port>/<target-host>/<path>

Examples:
    python proxy.py 300               # 5-minute cache on port 8888
    python proxy.py 3600 --port 9090  # 1-hour cache on port 9090
    python proxy.py 0                 # Disable caching (always fetch fresh)
"""

import argparse
import logging
import select
import signal
import socket
import sys
import time
from pathlib import Path
from typing import Optional

log = logging.getLogger("proxy")

_LISTEN_BACKLOG = 50
_RECV_SIZE = 4096


# ── Startup ──────────────────────────────────────────────────────────────────

def setup_logging(verbose: bool) -> None:
    logging.basicConfig(
        format="%(asctime)s %(levelname)-8s %(message)s",
        datefmt="%H:%M:%S",
        level=logging.DEBUG if verbose else logging.INFO,
    )


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        prog="proxy.py",
        description="Caching HTTP proxy server.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
examples:
  python proxy.py 60                      60-second cache on port 8888
  python proxy.py 3600 --port 9090        1-hour cache on port 9090
  python proxy.py 300 --upstream-port 8000
  python proxy.py 0                       disable caching

access sites via:
  http://localhost:8888/<hostname>/<path>
  e.g.  http://localhost:8888/example.com/index.html
        http://localhost:8888/localhost/mysite/
        """,
    )
    p.add_argument("stale_time", type=float,
                   help="Cache TTL in seconds (0 = disable caching)")
    p.add_argument("--port", type=int, default=8888,
                   help="Port to listen on (default: 8888)")
    p.add_argument("--upstream-port", type=int, default=80,
                   help="Port to connect to on upstream servers (default: 80)")
    p.add_argument("--host", default="127.0.0.1",
                   help="Address to bind to (default: 127.0.0.1). "
                        "Use 0.0.0.0 to accept connections from other machines.")
    p.add_argument("--cache-dir", default="cache",
                   help="Directory for cached responses (default: ./cache)")
    p.add_argument("--verbose", "-v", action="store_true",
                   help="Enable debug logging")
    return p.parse_args()


# ── HTTP parsing ─────────────────────────────────────────────────────────────

def parse_request(raw: str):
    """Parse 'GET /<host>/<path> HTTP/x.x' into (host, request_bytes, endpoint).

    Raises ValueError on malformed input.
    """
    try:
        first_line = raw.split("\r\n")[0]
        _method, path, *_ = first_line.split()
        segments = path.lstrip("/").split("/", 1)
        host = segments[0]
        endpoint = "/" + (segments[1] if len(segments) > 1 else "")
        if not host or ("." not in host and host != "localhost"):
            raise ValueError(f"invalid or missing host {host!r}")
        get_req = (
            b"GET " + endpoint.encode() +
            b" HTTP/1.0\r\nHost: " + host.encode() +
            b"\r\n\r\n"
        )
        return host, get_req, endpoint
    except (ValueError, IndexError) as exc:
        raise ValueError(f"malformed request: {exc}") from exc


# ── Cache helpers ─────────────────────────────────────────────────────────────

def _cache_key(host: str, endpoint: str) -> str:
    return (host + endpoint).replace("/", "_").replace(".", "_").replace(":", "_")


def cache_path(cache_dir: Path, host: str, endpoint: str) -> Path:
    return cache_dir / (_cache_key(host, endpoint) + ".cache")


def cache_valid(path: Path, stale_time: float) -> bool:
    if stale_time <= 0:
        return False
    try:
        return (time.time() - path.stat().st_mtime) < stale_time
    except FileNotFoundError:
        return False


def cache_store(path: Path, data: bytes) -> None:
    try:
        path.write_bytes(data)
        log.debug("stored  %s (%d B)", path.name, len(data))
    except OSError as exc:
        log.warning("cache write failed %s: %s", path, exc)


def cache_load(path: Path) -> Optional[bytes]:
    try:
        return path.read_bytes()
    except OSError as exc:
        log.warning("cache read failed %s: %s", path, exc)
        return None


# ── Socket helpers ────────────────────────────────────────────────────────────

def open_upstream(host: str, port: int) -> socket.socket:
    """Create a non-blocking TCP socket and begin connecting to host:port."""
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setblocking(False)
    try:
        sock.connect((host, port))
    except BlockingIOError:
        pass  # Expected with non-blocking connect — connection is in progress
    return sock


def _remove(sock: socket.socket, *lists) -> None:
    """Remove sock from each list if present (no error if missing)."""
    for lst in lists:
        try:
            lst.remove(sock)
        except ValueError:
            pass


def _try_send(sock: socket.socket, data: bytes,
              outputs: list, pending: dict) -> None:
    """Send data to sock. Store any unsent remainder in pending for later."""
    try:
        sent = sock.send(data)
        remainder = data[sent:]
        if remainder:
            pending[sock] = remainder
        else:
            _remove(sock, outputs)
            sock.close()
    except OSError:
        _remove(sock, outputs)
        try:
            sock.close()
        except OSError:
            pass


# ── Main proxy loop ──────────────────────────────────────────────────────────

def run(stale_time: float, listen_port: int, upstream_port: int,
        bind_host: str, cache_dir: Path) -> None:

    cache_dir.mkdir(parents=True, exist_ok=True)

    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.setblocking(False)
    srv.bind((bind_host, listen_port))
    srv.listen(_LISTEN_BACKLOG)

    log.info("Listening on %s:%d  |  cache TTL: %gs  |  cache dir: %s",
             bind_host, listen_port, stale_time, cache_dir.resolve())
    log.info("Access sites via  http://localhost:%d/<hostname>/<path>", listen_port)
    log.info("Press Ctrl-C to stop.")

    inputs:  list = [srv]
    outputs: list = []

    clients:   set = set()  # browser-side sockets
    upstreams: set = set()  # upstream-server sockets

    # Per-socket state
    partial_requests: dict = {}  # client   -> bytes  (accumulating browser request)
    upstream_req:     dict = {}  # upstream -> bytes  (GET request to send upstream)
    upstream_client:  dict = {}  # upstream -> client (which browser owns this conn)
    upstream_url:     dict = {}  # upstream -> (host, endpoint)
    upstream_data:    dict = {}  # upstream -> bytes  (accumulating upstream response)
    ready_resp:       dict = {}  # client   -> bytes  (full response ready to forward)
    cached_resp:      dict = {}  # client   -> bytes  (cache-hit response)
    pending:          dict = {}  # client   -> bytes  (partial send in progress)

    # ── Graceful shutdown ────────────────────────────────────────────────────
    alive = [True]

    def _stop(sig, _frame):
        log.info("Caught signal %d — shutting down…", sig)
        alive[0] = False

    signal.signal(signal.SIGINT, _stop)
    try:
        signal.signal(signal.SIGTERM, _stop)
    except (AttributeError, OSError):
        pass  # SIGTERM not available on all platforms

    def close(sock: socket.socket) -> None:
        """Remove sock from all tracking structures and close it."""
        _remove(sock, inputs, outputs)
        clients.discard(sock)
        upstreams.discard(sock)
        for d in (partial_requests, upstream_req, upstream_client,
                  upstream_url, upstream_data, ready_resp, cached_resp, pending):
            d.pop(sock, None)
        try:
            sock.close()
        except OSError:
            pass

    # ── Event loop ───────────────────────────────────────────────────────────
    while alive[0]:
        try:
            readable, writable, exceptional = select.select(
                inputs, outputs, inputs, 1.0
            )
        except (OSError, ValueError):
            break

        # ── Readable sockets ─────────────────────────────────────────────────
        for s in readable:

            # New browser connection
            if s is srv:
                conn, addr = srv.accept()
                conn.setblocking(False)
                inputs.append(conn)
                clients.add(conn)
                partial_requests[conn] = b""
                log.debug("connect  %s:%d", *addr)

            # Incoming data from a browser
            elif s in clients:
                chunk = s.recv(_RECV_SIZE)
                if not chunk:
                    close(s)
                    continue

                partial_requests[s] += chunk
                if b"\r\n\r\n" not in partial_requests[s]:
                    continue  # Request not yet complete

                raw = partial_requests.pop(s)
                _remove(s, inputs)
                clients.discard(s)

                try:
                    host, get_req, endpoint = parse_request(
                        raw.decode(errors="replace")
                    )
                except ValueError as exc:
                    log.warning("bad request: %s", exc)
                    close(s)
                    continue

                cpath = cache_path(cache_dir, host, endpoint)
                if cache_valid(cpath, stale_time):
                    data = cache_load(cpath)
                    if data:
                        log.info("HIT   %s%s", host, endpoint)
                        cached_resp[s] = data
                        outputs.append(s)
                        continue
                    # Cache file disappeared between check and read — fall through

                log.info("FETCH %s%s", host, endpoint)
                try:
                    up = open_upstream(host, upstream_port)
                except OSError as exc:
                    log.warning("connect failed %s:%d — %s", host, upstream_port, exc)
                    close(s)
                    continue

                inputs.append(up)
                outputs.append(up)
                upstreams.add(up)
                upstream_req[up]    = get_req
                upstream_client[up] = s
                upstream_url[up]    = (host, endpoint)
                upstream_data[up]   = b""

            # Incoming data from an upstream server
            else:
                chunk = s.recv(_RECV_SIZE)
                if chunk:
                    upstream_data[s] += chunk
                else:
                    # EOF — upstream has sent the full response
                    host, endpoint = upstream_url[s]
                    response = upstream_data[s]
                    client   = upstream_client[s]

                    cpath = cache_path(cache_dir, host, endpoint)
                    cache_store(cpath, response)
                    log.info("DONE  %s%s  (%d B)", host, endpoint, len(response))

                    ready_resp[client] = response
                    if client not in outputs:
                        outputs.append(client)

                    _remove(s, inputs, outputs)
                    upstreams.discard(s)
                    for d in (upstream_req, upstream_client,
                              upstream_url, upstream_data):
                        d.pop(s, None)
                    try:
                        s.close()
                    except OSError:
                        pass

        # ── Writable sockets ─────────────────────────────────────────────────
        for s in writable:

            # Non-blocking connect completed — send the GET request
            if s in upstreams:
                try:
                    s.sendall(upstream_req.pop(s))
                    _remove(s, outputs)
                    if s not in inputs:
                        inputs.append(s)
                except OSError as exc:
                    log.warning("upstream send failed: %s", exc)
                    client = upstream_client.get(s)
                    close(s)
                    if client:
                        close(client)

            # Resume a partial send to the browser
            elif s in pending:
                try:
                    sent = s.send(pending[s])
                    pending[s] = pending[s][sent:]
                    if not pending[s]:
                        pending.pop(s)
                        _remove(s, outputs)
                        s.close()
                except OSError:
                    close(s)

            # Send a cache-hit response to the browser
            elif s in cached_resp:
                _try_send(s, cached_resp.pop(s), outputs, pending)

            # Send a freshly fetched response to the browser
            elif s in ready_resp:
                _try_send(s, ready_resp.pop(s), outputs, pending)

            else:
                _remove(s, outputs)

        # ── Exceptional sockets ──────────────────────────────────────────────
        for s in exceptional:
            log.debug("socket error: %s", s)
            client = upstream_client.get(s)
            close(s)
            if client:
                close(client)

    log.info("Proxy stopped.")
    try:
        srv.close()
    except OSError:
        pass


# ── Entry point ───────────────────────────────────────────────────────────────

def main() -> None:
    args = parse_args()
    setup_logging(args.verbose)
    run(
        stale_time=args.stale_time,
        listen_port=args.port,
        upstream_port=args.upstream_port,
        bind_host=args.host,
        cache_dir=Path(args.cache_dir),
    )


if __name__ == "__main__":
    main()