requests patch/adapter to capture CONNECT response headers (Python 3)

requests_proxy_headers_3.py

# https://stackoverflow.com/questions/39068998/reading-connect-headers
# updated for Python 3
# WARNING: barely tested

import socket
import requests

from urllib3.connection import HTTPSConnection
from urllib3.connectionpool import HTTPSConnectionPool
from urllib3.poolmanager import ProxyManager

from requests.adapters import HTTPAdapter

_MAXLINE = 65536

class ProxyHeaderHTTPSConnection(HTTPSConnection):
    def __init__(self, *args, **kwargs):
        super(ProxyHeaderHTTPSConnection, self).__init__(*args, **kwargs)
        self._proxy_headers = []

    def _tunnel(self):
        def b(s): return s.encode()
        self.send(b("CONNECT %s:%d HTTP/1.0\r\n" % (self._tunnel_host, self._tunnel_port)))

        for header, value in self._tunnel_headers.items():
            self.send(b("%s: %s\r\n" % (header, value)))

        self.send(b("\r\n"))

        response = self.response_class(self.sock, method=self._method)
        version, code, message = response._read_status()

        if version == "HTTP/0.9":
            # HTTP/0.9 doesn't support the CONNECT verb, so if httplib has
            # concluded HTTP/0.9 is being used something has gone wrong.
            self.close()
            raise socket.error("Invalid response from tunnel request")

        if code != 200:
            self.close()
            raise socket.error("Tunnel connection failed: %d %s" % (code, message.strip()))

        self._proxy_headers = []

        while True:
            line = response.fp.readline(_MAXLINE + 1)

            if len(line) > _MAXLINE:
                raise Exception("header line too long")

            if not line or line == b'\r\n':
                break

            # The line is a header, save it
            if b':' in line:
                self._proxy_headers.append(tuple(v.strip() for v in line.decode().split(':', maxsplit=1)))

    def getresponse(self, *args, **kwargs):
        response = super(ProxyHeaderHTTPSConnection, self).getresponse(*args, **kwargs)
        response.msg._headers += self._proxy_headers

        return response


class ProxyHeaderHTTPSConnectionPool(HTTPSConnectionPool):
    ConnectionCls = ProxyHeaderHTTPSConnection


class ProxyHeaderProxyManager(ProxyManager):
    def _new_pool(self, scheme, host, port, request_context=None):
        assert scheme == 'https'

        return ProxyHeaderHTTPSConnectionPool(host, port, **self.connection_pool_kw)


class ProxyHeaderHTTPAdapter(HTTPAdapter):
    def proxy_manager_for(self, proxy, **proxy_kwargs):
        if proxy in self.proxy_manager:
            manager = self.proxy_manager[proxy]
        else:
            proxy_headers = self.proxy_headers(proxy)
            manager = self.proxy_manager[proxy] = ProxyHeaderProxyManager(
                proxy_url=proxy,
                proxy_headers=proxy_headers,
                num_pools=self._pool_connections,
                maxsize=self._pool_maxsize,
                block=self._pool_block,
                **proxy_kwargs)

        return manager

def mount_proxy_headers(session):
    session.mount('https://', ProxyHeaderHTTPAdapter())

def patch_http_adapter():
    HTTPAdapter.proxy_manager_for = ProxyHeaderHTTPAdapter.proxy_manager_for