#!/usr/bin/env python3
"""socks2http: a minimal HTTP forward proxy that tunnels to a SOCKS5 upstream.

Purpose: the Strata Docker build. The build container may not reach the public
internet directly, but may reach a SOCKS5 endpoint (e.g. socks5://myproxy:1080).
Standard clients in the build (apt, pip, Python's urllib) all natively honour an
HTTP/HTTPS *forward proxy* named in http_proxy / https_proxy - but they do NOT
natively speak a raw SOCKS5 URL. This process is that missing bridge:

    apt / pip / urllib  --http_proxy-->  socks2http (127.0.0.1:PORT)  --socks5-->  upstream

Run it on loopback during a `docker build` RUN that does network I/O:

    python3 socks2http.py --listen 127.0.0.1:12808 --upstream socks5://myproxy:1080 &
    export http_proxy=http://127.0.0.1:12808 https_proxy=http://127.0.0.1:12808
    ... run the network command ...
    kill %1

If --upstream is empty it runs as a plain TCP forward proxy (CONNECT still
tunnels, but the egress is a direct socket) - handy for testing / passthrough.

Standard library only. Binds to loopback by default; it does no authentication
of its own (it is only reachable on 127.0.0.1 inside the build container).
"""
from __future__ import annotations

import argparse
import socket
import socketserver
import sys
import threading


def parse_upstream(spec: str):
    """socks5://[user:pass@]host:port  -> (host, port, user, pass). '' -> None (passthrough)."""
    spec = (spec or "").strip()
    if not spec:
        return None
    if "://" not in spec:
        spec = "socks5://" + spec
    scheme, rest = spec.split("://", 1)
    if scheme not in ("socks5", "socks5h", "socks"):
        raise SystemExit(f"socks2http: unsupported upstream scheme '{scheme}' (use socks5://host:port)")
    user = pw = None
    if "@" in rest:
        auth, rest = rest.rsplit("@", 1)
        if ":" in auth:
            user, pw = auth.split(":", 1)
        else:
            user = auth
    host, _, port = rest.rpartition(":")
    if not host or not port.isdigit():
        raise SystemExit(f"socks2http: bad upstream '{spec}' (want socks5://host:port)")
    return host, int(port), user, pw


def socks5_connect(dest_host: str, dest_port: int, up) -> socket.socket:
    """Open a TCP socket to dest_host:dest_port through the SOCKS5 upstream (remote DNS)."""
    up_host, up_port, user, pw = up
    s = socket.create_connection((up_host, up_port), timeout=30)
    try:
        # greeting: VER 5, one method. 0x00 = no-auth, 0x02 = username/password.
        methods = b"\x02" if user else b"\x00"
        s.sendall(b"\x05\x01" + methods)
        ver, rep = s.recv(2)
        if ver != 5 or rep != 0:
            raise OSError(f"SOCKS5 greeting rejected: {ver!r} {rep!r}")
        if user:
            u = user.encode(); p = (pw or "").encode()
            s.sendall(b"\x01" + bytes([len(u)]) + u + bytes([len(p)]) + p)
            ver, rep = s.recv(2)
            if ver != 1 or rep != 0:
                raise OSError(f"SOCKS5 auth rejected: {ver!r} {rep!r}")
        # request: VER 5, CMD 1 (connect), RSV 0, ATYP 3 (domain, remote resolve)
        h = dest_host.encode("idna") if any(ord(c) > 127 for c in dest_host) else dest_host.encode()
        if len(h) > 255:
            raise OSError("dest hostname too long for SOCKS5")
        req = b"\x05\x01\x00\x03" + bytes([len(h)]) + h + dest_port.to_bytes(2, "big")
        s.sendall(req)
        resp = s.recv(4)
        if len(resp) < 4:
            raise OSError("short SOCKS5 response")
        if resp[1] != 0:
            raise OSError(f"SOCKS5 connect failed, rep={resp[1]}")
        atyp = resp[3]
        if atyp == 1:
            s.recv(4)
        elif atyp == 3:
            s.recv(1 + s.recv(1)[0])
        elif atyp == 4:
            s.recv(16)
        s.recv(2)  # bound port
        return s
    except Exception:
        s.close()
        raise


def _copy(a: socket.socket, b: socket.socket) -> None:
    """Copy from socket a to socket b until EOF/error. One direction only."""
    try:
        while True:
            data = a.recv(65536)
            if not data:
                break
            b.sendall(data)
    except OSError:
        pass
    finally:
        # Half-close the write side of b so the peer sees EOF, without tearing
        # down the socket the other direction is still using.
        try:
            b.shutdown(socket.SHUT_WR)
        except OSError:
            pass


def _relay(a: socket.socket, b: socket.socket) -> None:
    """Relay traffic between two sockets in both directions.

    Returns when both directions have reached EOF / error. Closes both sockets
    only after the relay is fully done (unlike a naive per-direction cleanup,
    which would kill the upstream before its response is read back).
    """
    t = threading.Thread(target=_copy, args=(b, a), daemon=True)
    t.start()
    _copy(a, b)
    t.join(timeout=30)
    for sk in (a, b):
        try:
            sk.shutdown(socket.SHUT_RDWR)
        except OSError:
            pass
        try:
            sk.close()
        except OSError:
            pass


class _Handler(socketserver.BaseRequestHandler):
    def handle(self) -> None:
        req = self.request
        up = self.server.upstream
        try:
            # Read the request head (first blank line).
            head = b""
            while b"\r\n\r\n" not in head:
                chunk = req.recv(65536)
                if not chunk:
                    return
                head += chunk
                if len(head) > 1 << 20:
                    return
            first = head.split(b"\r\n", 1)[0].decode("latin-1")
            parts = first.split()
            if len(parts) < 3:
                return

            if parts[0] == "CONNECT":
                # CONNECT host:port HTTP/1.1
                host, _, port = parts[1].rpartition(":")
                port = int(port or 443)
                if up:
                    egress = socks5_connect(host, port, up)
                else:
                    egress = socket.create_connection((host, port), timeout=30)
                req.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n")
                _relay(req, egress)
                return

            # Plain HTTP proxy: GET/POST/HEAD ... absolute-URI ...
            target = parts[1]
            if target.startswith("http://") or target.startswith("https://"):
                scheme, rest = target.split("://", 1)
                hostport, _, path = rest.partition("/")
                if ":" in hostport:
                    host, _, port = hostport.rpartition(":")
                    port = int(port)
                else:
                    host, port = hostport, (443 if scheme == "https" else 80)
                if up:
                    egress = socks5_connect(host, port, up)
                else:
                    egress = socket.create_connection((host, port), timeout=30)
                # Rewrite to origin-form and forward the rest of the head as-is.
                path = "/" + path
                rewritten = f"{parts[0]} {path} {parts[2]}\r\n".encode("latin-1")
                body = head.split(b"\r\n", 1)[1]
                egress.sendall(rewritten + body)
                _relay(req, egress)
                return

            self.request.sendall(b"HTTP/1.1 502 Bad Gateway\r\n\r\n")
        except (OSError, ValueError) as e:
            try:
                self.request.sendall(f"HTTP/1.1 502 Bad Gateway\r\n\r\n{e}\r\n".encode())
            except OSError:
                pass
            finally:
                try:
                    self.request.close()
                except OSError:
                    pass


class _Server(socketserver.ThreadingTCPServer):
    daemon_threads = True
    allow_reuse_address = True
    upstream = None


def main() -> int:
    ap = argparse.ArgumentParser(description="SOCKS5 -> HTTP forward proxy (loopback)")
    ap.add_argument("--listen", default="127.0.0.1:12808", help="bind addr:port (default 127.0.0.1:12808)")
    ap.add_argument("--upstream", default="",
                    help="socks5://host:port, or empty for a plain TCP forward proxy")
    a = ap.parse_args()
    host, _, port = a.listen.rpartition(":")
    host = host or "127.0.0.1"
    port = int(port)
    if not (0 < port < 65536):
        raise SystemExit(f"socks2http: --listen port {port} out of range (1-65535)")
    _Server.upstream = parse_upstream(a.upstream)
    srv = _Server((host, port), _Handler)
    sys.stderr.write(f"socks2http: listening on {host}:{port}, upstream="
                     f"{a.upstream or '(passthrough)'}\n")
    sys.stderr.flush()
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        pass
    return 0


if __name__ == "__main__":
    raise SystemExit(main())