outgoing_proxy.py

#!/usr/bin/env python

import logging
import socket
import urllib
import ipaddress

import tornado.ioloop
import tornado.iostream
import tornado.web
import tornado.httpclient
import tornado.httputil

logger = logging.getLogger('tornado_proxy')

def fetch_request(url, callback, **kwargs):
    req = tornado.httpclient.HTTPRequest(url, **kwargs)
    client = tornado.httpclient.AsyncHTTPClient()
    client.fetch(req, callback, raise_error=False)

def resolve_host(host):
    # Remove port from host if it exists
    parts = host.split(":")
    hostname = parts[0]
    try:
        host_ip = socket.gethostbyname(hostname)
        logger.debug('Resolved %s to %s', hostname, host_ip)
        ip_addr = ipaddress.IPv4Address(host_ip)
        if not ip_addr.is_global:
            return
        netloc = host_ip
        # If there was a port in the original hostname
        if len(parts) > 1:
            netloc += ':' + parts[1]
        return netloc
    except (socket.gaierror, UnicodeError, ipaddress.AddressValueError):
        logger.debug('Failed to resolve %s', host)
        return

class ProxyHandler(tornado.web.RequestHandler):
    SUPPORTED_METHODS = tornado.web.RequestHandler.SUPPORTED_METHODS + ('CONNECT',)

    def compute_etag(self):
        return None

    @tornado.web.asynchronous
    def get(self):
        logger.debug('Handle %s request to %s', self.request.method,
                     self.request.uri)

        def handle_response(response):
            if (response.error and not
                    isinstance(response.error, tornado.httpclient.HTTPError)):
                self.set_status(500)
                self.write('Internal server error:\n' + str(response.error))
            else:
                self.set_status(response.code, response.reason)
                self._headers = response.headers
                if response.body:
                    self.write(response.body)
            self.finish()

        body = self.request.body
        if not body:
            body = None

        if 'Proxy-Connection' in self.request.headers:
            del self.request.headers['Proxy-Connection']
        if 'Connection' in self.request.headers:
            del self.request.headers['Connection']

        parsed_url = urllib.parse.urlparse(self.request.uri)

        # The host header can be spoofed, so we get the host from the url itself.
        actual_host = parsed_url.netloc
        if actual_host != self.request.host:
            self.set_status(400, 'Invalid Host Header')
            self.finish()
            return

        resolved_netloc = resolve_host(self.request.host)
        if not resolved_netloc:
            self.set_status(404, 'Not Found')
            self.finish()
            return

        self.request.headers.add('Connection', 'Close')
        if 'Host' not in self.request.headers:
            self.request.headers.add('Host', self.request.host)

        final_url = urllib.parse.urlunparse((parsed_url.scheme, resolved_netloc,
                                             parsed_url.path, parsed_url.params,
                                             parsed_url.query, parsed_url.fragment))

        fetch_request(
            final_url, handle_response,
            method=self.request.method, body=body,
            headers=self.request.headers, follow_redirects=False,
            allow_nonstandard_methods=True, decompress_response=False)

    @tornado.web.asynchronous
    def head(self):
        return self.get()

    @tornado.web.asynchronous
    def post(self):
        return self.get()

    @tornado.web.asynchronous
    def delete(self):
        return self.get()

    @tornado.web.asynchronous
    def patch(self):
        return self.get()

    @tornado.web.asynchronous
    def put(self):
        return self.get()

    @tornado.web.asynchronous
    def options(self):
        return self.get()

    @tornado.web.asynchronous
    def connect(self):
        logger.debug('Start CONNECT to %s', self.request.uri)
        host, port = self.request.uri.split(':')
        client = self.request.connection.stream

        def read_from_client(data):
            upstream.write(data)

        def read_from_upstream(data):
            client.write(data)

        def client_close(data=None):
            if upstream.closed():
                return
            if data:
                upstream.write(data)
            upstream.close()

        def upstream_close(data=None):
            if client.closed():
                return
            if data:
                client.write(data)
            client.close()

        def start_tunnel():
            logger.debug('CONNECT tunnel established to %s', self.request.uri)
            client.read_until_close(client_close, read_from_client)
            upstream.read_until_close(upstream_close, read_from_upstream)
            client.write(b'HTTP/1.0 200 Connection established\r\n\r\n')

        host_ip = resolve_host(host)
        if not host_ip:
            self.set_status(404, 'Not Found')
            self.finish()
            return
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        upstream = tornado.iostream.IOStream(s)
        upstream.connect((host_ip, int(port)), start_tunnel)


def run_proxy(port):
    app = tornado.web.Application([
        (r'.*', ProxyHandler),
    ])
    app.listen(port)
    logger.debug('Serving on http://localhost:' + str(port))
    tornado.ioloop.IOLoop.current().start()

if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - pid:%(process)d - %(message)s')
    run_proxy(8899)