# -*- coding: utf-8 -*-
"""Local auth-free HTTP/CONNECT relay -> upstream proxy that requires credentials.

WHY: Chrome `--proxy-server=http://user:pass@host:port` fails with
ERR_NO_SUPPORTED_PROXIES (Chrome refuses credentials in the proxy URL). Point
Chrome/curl at this relay instead; it injects Proxy-Authorization upstream.
Also cleaner than CDP Fetch.enable(handleAuthRequests) with pattern '*' which
pauses every subresource and hangs the driver.

Usage:
  python auth_proxy_relay.py <local_port> <up_host> <up_port> <user> <pass>
  # then: chrome --proxy-server=http://127.0.0.1:<local_port>
  #       curl  -x http://127.0.0.1:<local_port> https://target/...

Verify egress before relying on it:
  curl -s -x http://127.0.0.1:<local_port> https://ipinfo.io/json
"""
import base64, select, socket, sys, threading


def handle(client, up_host, up_port, auth_header):
    up = None
    try:
        data = b""
        while b"\r\n\r\n" not in data:
            chunk = client.recv(65536)
            if not chunk:
                return
            data += chunk
        head, rest = data.split(b"\r\n\r\n", 1)
        lines = head.decode("latin1").split("\r\n")
        method, target = lines[0].split()[0], lines[0].split()[1]
        up = socket.create_connection((up_host, up_port), timeout=30)
        if method == "CONNECT":
            up.sendall(
                f"CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n"
                f"Proxy-Authorization: {auth_header}\r\n\r\n".encode()
            )
            resp = b""
            while b"\r\n\r\n" not in resp:
                c = up.recv(4096)
                if not c:
                    break
                resp += c
            first = resp.split(b"\r\n", 1)[0].decode("latin1")
            # forward upstream status line verbatim to the client
            client.sendall(b"HTTP/1.1 " + first.split(" ", 1)[1].encode() + b"\r\n\r\n")
            tail = resp.split(b"\r\n\r\n", 1)
            if len(tail) > 1 and tail[1]:
                client.sendall(tail[1])
        else:
            new_lines = [lines[0]]
            has_auth = False
            for l in lines[1:]:
                low = l.lower()
                if low.startswith("proxy-authorization"):
                    has_auth = True
                    continue
                if low.startswith("proxy-connection"):
                    continue
                new_lines.append(l)
            if not has_auth:
                new_lines.append(f"Proxy-Authorization: {auth_header}")
            up.sendall(("\r\n".join(new_lines) + "\r\n\r\n").encode() + rest)
        # bidirectional tunnel
        socks = [client, up]
        while True:
            r, _, x = select.select(socks, [], socks, 60)
            if x or not r:
                break
            for s in r:
                d = s.recv(65536)
                if not d:
                    return
                (up if s is client else client).sendall(d)
    except Exception as e:
        try:
            client.sendall(b"HTTP/1.1 502 Bad Gateway\r\n\r\n" + str(e).encode()[:200])
        except Exception:
            pass
    finally:
        for s in (client, up):
            try:
                if s:
                    s.close()
            except Exception:
                pass


def main():
    if len(sys.argv) != 6:
        print(__doc__)
        sys.exit(2)
    port, up_host, up_port, user, pw = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4], sys.argv[5]
    auth = b"Basic " + base64.b64encode(f"{user}:{pw}".encode())
    srv = socket.socket()
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(("127.0.0.1", int(port)))
    srv.listen(128)
    print(f"relay listening 127.0.0.1:{port} -> {up_host}:{up_port}", flush=True)
    while True:
        cl, _ = srv.accept()
        threading.Thread(target=handle, args=(cl, up_host, up_port, auth.decode()), daemon=True).start()


if __name__ == "__main__":
    main()