proxy.py

"""Async HTTP→SOCKS5 proxy with SSH tunnel management."""
import asyncio
import logging
import os
import socket
import struct
import subprocess
import threading
import time
from collections import deque

from stats import Stats

logger = logging.getLogger("magic-proxy.proxy")

SOCKS5_TIMEOUT = 15            # seconds for SOCKS5 handshake
RELAY_BUF = 65536              # bytes per read
DRAIN_THRESHOLD = 256 * 1024   # bytes between explicit drains
PORT_PROBE_TIMEOUT = 0.1       # seconds for connecting-state probe


async def socks5_connect(host: str, port: int, socks_addr: str):
    """Establish a SOCKS5 connection (no auth) through the tunnel."""
    socks_host, socks_port = socks_addr.split(":")
    reader, writer = await asyncio.wait_for(
        asyncio.open_connection(socks_host, int(socks_port)),
        timeout=SOCKS5_TIMEOUT,
    )

    try:
        async def handshake():
            writer.write(b"\x05\x01\x00")
            await writer.drain()
            resp = await reader.readexactly(2)
            if resp != b"\x05\x00":
                raise RuntimeError(f"SOCKS5 handshake rejected: {resp.hex()}")

            try:
                ip = socket.inet_aton(host)
                addr_type = b"\x01"
                addr = ip
            except OSError:
                addr_type = b"\x03"
                addr = len(host).to_bytes(1, "big") + host.encode()

            writer.write(b"\x05\x01\x00" + addr_type + addr + struct.pack("!H", port))
            await writer.drain()

            header = await reader.readexactly(4)
            if header[0] != 0x05 or header[1] != 0x00:
                raise RuntimeError(f"SOCKS5 connect failed: status={header[1]}")

            atyp = header[3]
            if atyp == 0x01:
                await reader.readexactly(6)
            elif atyp == 0x03:
                length = (await reader.readexactly(1))[0]
                await reader.readexactly(length + 2)
            elif atyp == 0x04:
                await reader.readexactly(18)
            else:
                raise RuntimeError(f"SOCKS5 unknown address type: {atyp}")

        await asyncio.wait_for(handshake(), timeout=SOCKS5_TIMEOUT)
        return reader, writer
    except Exception:
        writer.close()
        try:
            await writer.wait_closed()
        except Exception:
            pass
        raise


async def bidirectional_relay(ra, wa, rb, wb, stats):
    """Half-close-aware relay: each direction drains independently, then EOFs."""
    async def relay(src, dst, record):
        bytes_since_drain = 0
        try:
            while True:
                data = await src.read(RELAY_BUF)
                if not data:
                    break
                dst.write(data)
                record(len(data))
                bytes_since_drain += len(data)
                if bytes_since_drain >= DRAIN_THRESHOLD:
                    await dst.drain()
                    bytes_since_drain = 0
            try:
                await dst.drain()
            except Exception:
                pass
        except (ConnectionError, OSError, asyncio.CancelledError):
            pass
        finally:
            try:
                if dst.can_write_eof():
                    dst.write_eof()
            except Exception:
                pass

    await asyncio.gather(
        relay(ra, wb, stats.record_up),
        relay(rb, wa, stats.record_down),
        return_exceptions=True,
    )
    for w in (wa, wb):
        try:
            w.close()
        except Exception:
            pass


async def handle_connect(client_reader, client_writer, host, port, socks_addr, stats):
    """Handle HTTP CONNECT (HTTPS) tunneling."""
    stats.inc_connections()
    try:
        while True:
            line = await client_reader.readline()
            if line == b"\r\n" or line == b"":
                break

        try:
            remote_reader, remote_writer = await socks5_connect(host, port, socks_addr)
        except Exception as e:
            logger.error("SOCKS5 connect to %s:%s failed: %s", host, port, e)
            client_writer.write(b"HTTP/1.1 502 Bad Gateway\r\n\r\n")
            await client_writer.drain()
            client_writer.close()
            return

        client_writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
        await client_writer.drain()
        await bidirectional_relay(client_reader, client_writer, remote_reader, remote_writer, stats)
    finally:
        stats.dec_connections()


async def handle_http(client_reader, client_writer, request_line, socks_addr, stats):
    """Handle plain HTTP request via SOCKS5."""
    stats.inc_connections()
    try:
        method, url, _ = request_line.decode().split(maxsplit=2)

        if url.startswith("http://"):
            url = url[7:]
        host, _, rest = url.partition("/")
        if ":" in host:
            host, port_str = host.split(":", 1)
            port = int(port_str)
        else:
            port = 80

        try:
            remote_reader, remote_writer = await socks5_connect(host, port, socks_addr)
        except Exception as e:
            logger.error("SOCKS5 connect to %s:%s failed: %s", host, port, e)
            client_writer.write(b"HTTP/1.1 502 Bad Gateway\r\n\r\n")
            await client_writer.drain()
            client_writer.close()
            return

        path = "/" + rest if rest else "/"
        remote_writer.write(f"{method} {path} HTTP/1.1\r\n".encode())

        while True:
            line = await client_reader.readline()
            if line == b"\r\n" or line == b"":
                break
            remote_writer.write(line)

        remote_writer.write(b"\r\n")
        await remote_writer.drain()
        await bidirectional_relay(client_reader, client_writer, remote_reader, remote_writer, stats)
    finally:
        stats.dec_connections()


async def handle_client(client_reader, client_writer, socks_addr, stats):
    """Process one HTTP proxy client."""
    try:
        request_line = await client_reader.readline()
        if not request_line:
            client_writer.close()
            return

        parts = request_line.decode(errors="replace").split()
        if len(parts) < 2:
            client_writer.close()
            return
        method, target = parts[0], parts[1]

        if method == "CONNECT":
            host, _, port_str = target.partition(":")
            if not port_str:
                client_writer.close()
                return
            await handle_connect(client_reader, client_writer, host, int(port_str), socks_addr, stats)
        else:
            await handle_http(client_reader, client_writer, request_line, socks_addr, stats)
    except asyncio.CancelledError:
        raise
    except (ConnectionResetError, BrokenPipeError, asyncio.IncompleteReadError) as e:
        # Client closed mid-stream — normal browser/curl behavior, not an error.
        logger.debug("client disconnected: %s", e)
        try:
            client_writer.close()
        except Exception:
            pass
    except Exception:
        logger.exception("Error handling client")
        try:
            client_writer.close()
        except Exception:
            pass


class SSHMonitor:
    """Manages an SSH SOCKS5 tunnel subprocess."""

    def __init__(self):
        self.process = None
        self._status = "stopped"
        self._error_msg = ""
        self._cmd_str = ""
        self._log_lines = deque(maxlen=100)
        self._log_lock = threading.Lock()
        self._start_time = 0
        self._log_thread = None
        self._current_name = ""

    @property
    def status(self):
        return self._status

    @property
    def error_msg(self):
        return self._error_msg

    @property
    def cmd_str(self):
        return self._cmd_str

    @property
    def log(self) -> str:
        with self._log_lock:
            lines = list(self._log_lines)
        return "\n".join(lines[-5:])

    @property
    def elapsed(self) -> float:
        if self._start_time and self._status in ("connecting", "connected"):
            return time.monotonic() - self._start_time
        return 0

    @property
    def current_name(self) -> str:
        return self._current_name

    def start(self, tunnel: dict, socks5_port: int, password: str = ""):
        """Start the SSH tunnel subprocess for the given tunnel config."""
        self.stop()

        user = tunnel.get("ssh_user", "")
        host = tunnel["ssh_host"]
        port = str(tunnel.get("ssh_port", 22))
        auth_type = tunnel.get("auth_type", "key")
        compression = tunnel.get("ssh_compression", True)

        destination = f"{user}@{host}" if user else host

        ssh_args = [
            "-D", str(socks5_port),
            "-N",
            "-o", "ExitOnForwardFailure=yes",
            "-o", "StrictHostKeyChecking=accept-new",
        ]
        if compression:
            ssh_args.append("-C")
        ssh_args.extend(["-p", port, destination])

        # sshpass via fd avoids leaking the password into `ps`.
        pass_fds = ()
        r_fd = None
        if auth_type == "password":
            r_fd, w_fd = os.pipe()
            try:
                os.write(w_fd, (password + "\n").encode())
            finally:
                os.close(w_fd)
            cmd = ["sshpass", "-d", str(r_fd), "ssh"] + ssh_args
            pass_fds = (r_fd,)
            display_cmd = ["sshpass", "-d", "***", "ssh"] + ssh_args
        else:
            key = tunnel.get("ssh_key", "")
            cmd = ["ssh", "-i", key] + ssh_args
            display_cmd = cmd

        self._cmd_str = " ".join(display_cmd)
        self._current_name = tunnel.get("name", destination)
        self._status = "connecting"
        self._error_msg = ""
        with self._log_lock:
            self._log_lines.clear()
        self._start_time = time.monotonic()

        try:
            self.process = subprocess.Popen(
                cmd,
                stdin=subprocess.DEVNULL,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
                pass_fds=pass_fds,
            )
        finally:
            if r_fd is not None:
                try:
                    os.close(r_fd)
                except OSError:
                    pass

        self._log_thread = threading.Thread(target=self._read_stderr, daemon=True)
        self._log_thread.start()
        logger.info("SSH starting: %s", self._cmd_str)

    def _read_stderr(self):
        # Owns the stderr pipe: closes it in finally so the fd always releases
        # (on EOF when ssh exits, or on exception). This close must NOT happen on
        # the main thread — BufferedReader.close() takes the buffer lock this
        # loop holds while blocked in read(), which deadlocked _quit.
        try:
            for line in self.process.stderr:
                decoded = line.decode(errors="replace").rstrip()
                if decoded:
                    with self._log_lock:
                        self._log_lines.append(decoded)
        except Exception:
            logger.exception("stderr reader crashed")
        finally:
            try:
                self.process.stderr.close()
            except Exception:
                pass

    def stop(self, blocking=True):
        """Stop the SSH process.

        Order matters: terminate() FIRST, then wait. We never close stderr here
        — the reader thread owns that pipe and closes it on EOF (see
        _read_stderr). Calling stderr.close() on this thread deadlocks:
        BufferedReader.close() blocks on the same buffer lock the reader holds
        while blocked in read(). That deadlock, not process.wait(), is the real
        "click Quit, app hangs" bug.

        blocking=False (quit path) sends SIGTERM and returns immediately; the
        pipe and reader thread are reclaimed when the process exits.
        """
        if self.process is None:
            return
        self.process.terminate()
        if not blocking:
            # Drop the handle; the OS process exits on the already-delivered SIGTERM.
            self.process = None
            self._status = "stopped"
            return
        try:
            self.process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            self.process.kill()
        self._wait_log_thread()
        self.process = None
        self._status = "stopped"

    def _wait_log_thread(self):
        if self._log_thread and self._log_thread.is_alive():
            self._log_thread.join(timeout=1)

    def check(self, socks5_port: int):
        """Poll the subprocess; only probe the SOCKS5 port while connecting."""
        if self.process is None:
            return

        ret = self.process.poll()
        if ret is not None:
            self._status = "error"
            self._wait_log_thread()
            with self._log_lock:
                lines = list(self._log_lines)
            self._error_msg = "\n".join(lines) or f"SSH exited with code {ret}"
            logger.warning("SSH exited (%s): %s", ret, self._error_msg)
            self.process = None
            return

        # Once we're connected the SSH process death is the only thing left to
        # observe; skipping the probe keeps the UI thread from blocking on
        # connect() every second.
        if self._status != "connecting":
            return

        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            s.settimeout(PORT_PROBE_TIMEOUT)
            s.connect(("127.0.0.1", socks5_port))
            self._status = "connected"
            logger.info("SSH tunnel ready on port %d", socks5_port)
        except (ConnectionRefusedError, OSError):
            pass
        finally:
            s.close()


async def run_proxy(config: dict, stats: Stats, control: dict):
    """Start the HTTP→SOCKS5 proxy server (runs in background thread's asyncio loop)."""
    socks_addr = f"127.0.0.1:{config['socks5_port']}"
    listen_host, listen_port = config["http_listen"].rsplit(":", 1)
    listen_port = int(listen_port)

    server = await asyncio.start_server(
        lambda r, w: handle_client(r, w, socks_addr, stats),
        listen_host, listen_port,
    )
    control["server"] = server
    logger.info("HTTP proxy listening on %s:%d → SOCKS5 %s", listen_host, listen_port, socks_addr)
    async with server:
        await server.serve_forever()