tornado http proxy

gistfile1.py

#!/usr/bin/env python
#coding:utf-8

import sys, os, re
import logging
import socket
import urlparse
from tornado import iostream, ioloop
from tornado.httpserver import HTTPServer
from tornado import httputil

logging.basicConfig(level=logging.INFO, format='%(levelname)s - - %(asctime)s %(message)s', datefmt='[%d/%b/%Y %H:%M:%S]')

class LocalProxyHandler(object):

    def __init__(self, request):
        self.request = request
        scheme, netloc, path, params, query, fragment = urlparse.urlparse(self.request.uri, 'http')
        self.remote_host, self.remote_port = self.resolve_netloc(netloc)
        self.remote_path = urlparse.urlunparse((scheme, self.remote_host + ('' if self.remote_host == 80 else ':%d' % self.remote_port), path, params, query, ''))
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        self.remote_stream = iostream.IOStream(sock)
        self.remote_stream.connect((self.remote_host, self.remote_port), self._on_remote_connect)

    def resolve_netloc(self, netloc):
        try:
            host, _, port = netloc.rpartition(':')
            return host, int(port)
        except ValueError:
            return netloc, 80

    def _on_remote_connect(self):
        logging.info('remote address (%r, %r) connected.', self.remote_host, self.remote_port)
        data ='%s %s %s\r\n'  % (self.request.method, self.remote_path, self.request.version)
        self.request.headers['Connection'] = 'close'
        self.request.headers['Proxy-Connection'] = 'close'
        data += ''.join('%s: %s\r\n' % (k, v) for k,v in self.request.headers.get_all())
        data += '\r\n'
        self.remote_stream.write(data, self._on_remote_write)

    def _on_remote_write(self):
        logging.info('remote address (%r, %r) send data end.', self.remote_host, self.remote_port)
        self.remote_stream.read_until('\r\n\r\n', self._on_remote_read_headers)

    def _on_remote_read_headers(self, data):
        logging.info('remote address (%r, %r) read headers end.', self.remote_host, self.remote_port)
        eol = data.find('\r\n')
        self.remote_headers = httputil.HTTPHeaders.parse(data[eol:])
        self.request.connection.write(data, self._on_local_write_headers)

    def _on_local_write_headers(self):
        logging.info('remote address (%r, %r) to local write headers end.', self.remote_host, self.remote_port)
        if self.remote_headers.get('connection') == 'close':
            self.remote_stream.read_until_close(self._on_remote_read_body)
        elif 'content-length' in self.remote_headers:
            self.remote_stream.read_bytes(int(self.remote_headers['content-length']), self._on_remote_read_body)
        else:
            self.remote_stream.read_until_close(self._on_remote_read_body)

    def _on_remote_read_body(self, data):
        logging.info('remote address (%r, %r) read finish.', self.remote_host, self.remote_port)
        self.request.write(data)
        self.request.finish()


def main():
    http_server = HTTPServer(LocalProxyHandler, no_keep_alive=True)
    http_server.listen(8888)
    ioloop.IOLoop.instance().start()

if __name__ == '__main__':
    main()