#!/usr/bin/env python3
"""Local HTTP forward proxy that injects Proxy-Authorization for upstream auth.
Use when DrissionPage/Playwright doesn't support authenticated proxies.

Usage:
  1. Edit UPSTREAM, AUTH_USER, AUTH_PASS below
  2. Run: python proxy_forwarder.py
  3. Set browser_proxy to http://127.0.0.1:9128
"""

import http.server
import urllib.request
import base64
import socket
import select

UPSTREAM = 'residential-x.bpproxy.at:1000'
AUTH_USER = b'bpuser-XXX'
AUTH_PASS = b'YYYY'
AUTH = base64.b64encode(AUTH_USER + b':' + AUTH_PASS).decode()
LOCAL = 9128


class Proxy(http.server.BaseHTTPRequestHandler):
    timeout = 30

    def do_GET(self):
        self._forward('GET')

    def do_POST(self):
        self._forward('POST')

    def do_HEAD(self):
        self._forward('HEAD')

    def do_PUT(self):
        self._forward('PUT')

    def _forward(self, method):
        try:
            url = self.path
            req = urllib.request.Request(url, method=method)
            req.add_header('Proxy-Authorization', f'Basic {AUTH}')
            for h, v in self.headers.items():
                if h.lower() not in ('host', 'proxy-connection', 'proxy-authorization'):
                    req.add_header(h, v)
            if method in ('POST', 'PUT'):
                body = self.rfile.read(int(self.headers.get('Content-Length', 0)))
                req.data = body
            req.set_proxy(UPSTREAM, 'http')
            resp = urllib.request.urlopen(req, timeout=30)
            self.send_response(resp.status)
            for h, v in resp.getheaders():
                self.send_header(h, v)
            self.end_headers()
            self.wfile.write(resp.read())
        except Exception as e:
            self.send_error(502, str(e))

    def do_CONNECT(self):
        host, port = self.path.split(':')
        port = int(port)
        try:
            uphost, upport = UPSTREAM.split(':')
            remote = socket.create_connection((uphost, int(upport)), timeout=10)
            connect_req = (
                f'CONNECT {host}:{port} HTTP/1.1\r\n'
                f'Host: {host}:{port}\r\n'
                f'Proxy-Authorization: Basic {AUTH}\r\n\r\n'
            )
            remote.send(connect_req.encode())
            resp = b''
            while b'\r\n\r\n' not in resp:
                resp += remote.recv(4096)
            if b'200' not in resp.split(b'\r\n')[0]:
                self.send_error(502, 'CONNECT failed')
                remote.close()
                return
            self.send_response(200, 'Connection Established')
            self.end_headers()
            self.connection.settimeout(30)
            remote.settimeout(30)
            socks = [self.connection, remote]
            while True:
                r, _, _ = select.select(socks, [], [], 30)
                if not r:
                    break
                for s in r:
                    try:
                        d = s.recv(8192)
                        if not d:
                            return
                        dst = remote if s is self.connection else self.connection
                        dst.sendall(d)
                    except Exception:
                        return
        except Exception as e:
            self.send_error(502, str(e))


if __name__ == '__main__':
    print(f'[*] Proxy Forwarder: 127.0.0.1:{LOCAL} -> {UPSTREAM}', flush=True)
    httpd = http.server.HTTPServer(('127.0.0.1', LOCAL), Proxy)
    httpd.serve_forever()