The simple HTTP proxy server

simpleproxy.py

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

import SocketServer
import BaseHTTPServer
import httplib, urlparse
import socket
import StringIO

import time
import logging
perflog = logging.getLogger(__name__)

class HTTPHeaderParser(BaseHTTPServer.BaseHTTPRequestHandler):
    """Parse http header and check it valid or invalid.

    @see http://stackoverflow.com/questions/4685217/parse-raw-http-headers
    """
    def __init__(self, request_text):
        self.rfile = StringIO.StringIO(request_text)
        self.raw_requestline = self.rfile.readline()
        self.error_code = self.error_message = None
        self.parse_request()

    def send_error(self, code, message):
        self.error_code = code
        self.error_message = message

class ProxyRequestHandler(SocketServer.StreamRequestHandler):
    """Handle proxy request

    Receive, parse, forward, and response to HTTP proxy request.
    """
    def __send_error(self, code):
        """Send error and close connection"""
        msg = "HTTP/1.1 {0} {1}\r\n\r\n".format(code,
                                                httplib.responses[code])
        self.request.send(msg)
        self.request.close()

        return

    def handle(self):
        """handle HTTP request"""
        ### Step 1: Receive request from client
        ## Parse header
        header_string = ""
        currentline = self.rfile.readline()
        while currentline.strip() != "":
            header_string += currentline
            currentline = self.rfile.readline()

        # validation
        header = HTTPHeaderParser(header_string)
        if header.error_code != None:
            self.__send_error(header.error_code)
            return
        if header.command not in ["GET", "POST"]:
            # Not implemented commands
            self.__send_error(501)
            return

        ## Get message body
        body = None
        if header.command == "POST":
            body = self.rfile.read(int(header.headers.dict['content-length']))

        ### Step 2: Forward request to server and receive responce
        location = urlparse.urlparse(header.path)
        connection = httplib.HTTPConnection(location.netloc)
        url = location.path + (""
                               if location.query == ""
                               else "?" + location.query)
        try:
            t1 = time.time()            # performance log 1
            connection.request(header.command, url, body, header.headers.dict)
            response = connection.getresponse()
        except socket.gaierror:
            # faild to resolve the name
            self.__send_error(502)
            return

        t2 = time.time()                # performace log 2

        ### Step 3: Send response to client
        # build header
        response_header_string = ""
        if response.version == 10:
            response_header_string += "HTTP/1.0 "
        else:
            response_header_string += "HTTP/1.1 "
        response_header_string += "{0} {1} \r\n".format(response.status,
                                                        response.reason)
        for h, v in response.getheaders():
            # FIXME: Support chunked response
            if h in ["transfer-encoding",]:
                continue
            response_header_string += "{0}: {1}\r\n".format(h, v)
        response_header_string += "\r\n"

        # send header and body
        self.request.send(response_header_string)
        self.request.send(response.read())

        # FIXME: support persistent connection
        self.request.close()

        # commit performace log
        perflog.info("{0} {1}".format(header.path, t2 - t1))

def parse_args():
    """Parse arguments"""
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--port", default=8080)
    args = parser.parse_args()

    return args

def main():
    """"""
    ## Parse arguments
    args = parse_args()

    # set debug level
    logging.basicConfig(format='%(message)s',
                        level = logging.INFO)

    ## Start proxy server service
    server = SocketServer.ThreadingTCPServer(("", args.port),
                                             ProxyRequestHandler)

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass

    server.server_close()

if __name__ == '__main__':
    main()