#!/usr/bin/env python3
"""HTTP (CONNECT + plain) proxy that forwards everything to a SOCKS5 upstream.

Why: a lot of tools/extensions only speak HTTP(S) proxies, not SOCKS. Run
this once and point them at it -- it bridges to your local SSH dynamic SOCKS
tunnel, doing remote DNS (socks5h) on the upstream side.

    +-----------+        HTTP/HTTPS         +-----------+      SOCKS5h       +-----------+
    |  tools    |  -->  127.0.0.1:18080  --> |  this     |  --> 127.0.0.1:12026 --> | ssh -D    |
    +-----------+                          |  bridge   |                     +-----------+
                                          +-----------+

Usage:
    .venv/bin/python http2socks.py                 # default 127.0.0.1:18080
    .venv/bin/python http2socks.py --port 8080      # custom port
    BIND=0.0.0.0 .venv/bin/python http2socks.py     # bind all interfaces
    UPSTREAM=socks5h://127.0.0.1:18080 .venv/bin/python http2socks.py

Env vars (override defaults):
    BIND        listen address            (default 127.0.0.1)
    PORT        listen port               (default 18080)
    UPSTREAM    socks5h://host:port        (default socks5h://127.0.0.1:12026)

Then set for your tools:
    HTTP_PROXY=http://127.0.0.1:18080
    HTTPS_PROXY=http://127.0.0.1:18080
    ALL_PROXY=http://127.0.0.1:18080
"""

import argparse
import os
import socket
import socketserver
import threading
import urllib.parse

import socks  # PySocks

# --- upstream config --------------------------------------------------------

def parse_upstream(s: str):
    u = urllib.parse.urlparse(s)
    scheme = (u.scheme or "socks5h").lower()
    if scheme not in ("socks5h", "socks5"):
        raise SystemExit(f"unsupported upstream scheme: {scheme!r} (use socks5h://)")
    # socks5h => PROXY_TYPE_SOCKS5 with remote DNS (rdns=True)
    host = u.hostname or "127.0.0.1"
    port = u.port or 12026
    return host, port


UPSTREAM_DEFAULT = os.environ.get("UPSTREAM", "socks5h://127.0.0.1:12026")
UP_HOST, UP_PORT = parse_upstream(UPSTREAM_DEFAULT)


def make_upstream_socket() -> socks.socksocket:
    s = socks.socksocket()
    # PROXY_TYPE_SOCKS5 + rdns=True => socks5h (DNS resolved by the SOCKS server)
    s.set_proxy(socks.SOCKS5, UP_HOST, UP_PORT, rdns=True)
    return s


# --- proxy server -----------------------------------------------------------

def pipe(src: socket.socket, dst: socket.socket):
    try:
        while True:
            data = src.recv(65536)
            if not data:
                break
            dst.sendall(data)
    except OSError:
        pass
    finally:
        for s in (src, dst):
            try:
                s.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass
            try:
                s.close()
            except OSError:
                pass


class ProxyHandler(socketserver.BaseRequestHandler):
    timeout = 120

    def handle(self):
        req = self.request
        # read the request line + headers
        buf = b""
        while b"\r\n\r\n" not in buf:
            chunk = req.recv(4096)
            if not chunk:
                return
            buf += chunk
            if len(buf) > 1 << 16:  # 64k guard against malformed clients
                req.sendall(b"HTTP/1.1 400 Bad Request\r\n\r\n")
                req.close()
                return

        head, _, rest = buf.partition(b"\r\n\r\n")
        lines = head.split(b"\r\n")
        request_line = lines[0].decode("latin-1")
        try:
            method, target, _ver = request_line.split(" ", 2)
        except ValueError:
            req.sendall(b"HTTP/1.1 400 Bad Request\r\n\r\n")
            req.close()
            return

        upstream = make_upstream_socket()

        if method == "CONNECT":
            # target is host:port
            host, _, port = target.rpartition(":")
            if not host or not port:
                req.sendall(b"HTTP/1.1 400 Bad Request\r\n\r\n")
                req.close()
                return
            try:
                upstream.connect((host, int(port)))
            except Exception:
                req.sendall(b"HTTP/1.1 502 Bad Gateway\r\n\r\n")
                req.close()
                upstream.close()
                return
            req.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n")
            # any bytes already read past the CONNECT line go upstream
            if rest:
                upstream.sendall(rest)
            threading.Thread(target=pipe, args=(req, upstream), daemon=True).start()
            pipe(upstream, req)

        else:
            # plain HTTP request; target is a full URL
            parsed = urllib.parse.urlparse(target)
            host = parsed.hostname
            port = parsed.port or (443 if parsed.scheme == "https" else 80)
            if not host:
                req.sendall(b"HTTP/1.1 400 Bad Request\r\n\r\n")
                req.close()
                upstream.close()
                return
            try:
                upstream.connect((host, port))
            except Exception:
                req.sendall(b"HTTP/1.1 502 Bad Gateway\r\n\r\n")
                req.close()
                upstream.close()
                return
            # forward original request verbatim
            upstream.sendall(buf)
            threading.Thread(target=pipe, args=(req, upstream), daemon=True).start()
            pipe(upstream, req)


class ThreadingTCPServer(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True


def main():
    ap = argparse.ArgumentParser(description="HTTP proxy bridging to a SOCKS5h upstream")
    ap.add_argument("--bind", default=os.environ.get("BIND", "127.0.0.1"))
    ap.add_argument("--port", type=int, default=int(os.environ.get("PORT", "18080")))
    args = ap.parse_args()

    print(f"http2socks: listening on http://{args.bind}:{args.port}  ->  "
          f"socks5h://{UP_HOST}:{UP_PORT} (remote DNS)", flush=True)
    with ThreadingTCPServer((args.bind, args.port), ProxyHandler) as srv:
        try:
            srv.serve_forever()
        except KeyboardInterrupt:
            print("\nhttp2socks: shutting down", flush=True)


if __name__ == "__main__":
    main()