simple_http_proxy_server

simple_http_proxy_server.py

# -*- coding: utf-8 -*-
# 15/8/21
# create by: snower

import sys
import time
import sevent
import logging
import traceback

def format_data_len(date_len):
    if date_len < 1024:
        return "%dB" % date_len
    elif date_len < 1024*1024:
        return "%.3fK" % (date_len/1024.0)
    elif date_len < 1024*1024*1024:
        return "%.3fM" % (date_len/(1024.0*1024.0))

class ProtocolParseEndError(Exception):
    def __init__(self, data, *args, **kwargs):
        super(ProtocolParseEndError, self).__init__(*args, **kwargs)
        self.data = data

class HttpProtocol(object):
    def __init__(self, request):
        self.request = request
        self.remote_addr = ''
        self.remote_port = 0
        self.method = None

    def get_method(self, data):
        index = data.find(" ")
        self.method = data[:index]
        return data[index+1:]

    def parse_addr(self, addr_info):
        addr_info = addr_info.split(":")
        self.remote_addr = addr_info[0]
        if len(addr_info) == 2:
            self.remote_port = int(addr_info[1])
        else:
            self.remote_port = 80
        if not self.remote_addr or not self.remote_port:
            raise Exception(addr_info)

    def parse_http(self, data):
        data = data[7:]
        index = data.find("/")
        self.parse_addr(data[:index])
        raise ProtocolParseEndError("".join([self.method, ' ', data[index:]]))

    def parse_https(self, data):
        index = data.find(" ")
        self.parse_addr(data[:index])
        self.request.write("HTTP/1.1 200 Connection Established\r\n\r\n")
        raise ProtocolParseEndError('')

    def parse(self, data):
        if self.method is None:
            data = self.get_method(data)
        if self.method.lower() == "connect":
            self.parse_https(data)
        elif self.method.lower() in ("get", "post", "put", "options", "head", "delete", "patch"):
            self.parse_http(data)
        else:
            raise Exception("unknown method %s" % self.method)

class Response(object):
    def __init__(self, request):
        self.conn = sevent.tcp.Socket()
        self.request = request
        self.is_connected = False
        self.buffer = None
        self.time = time.time()
        self.data_len = 0
        self.rbuffer, self.wbuffer = self.conn.buffer

        self.conn.on('connect', self.on_connect)
        self.conn.on('data', self.on_data)
        self.conn.on('close', self.on_close)
        self.conn.connect((self.request.protocol.remote_addr, self.request.protocol.remote_port), 30)

    def on_connect(self, s):
        self.is_connected = True
        self.wbuffer.link(self.request.rbuffer)
        self.request.wbuffer.link(self.rbuffer)
        if self.buffer:
            self.data_len += len(self.buffer)
            self.write(self.buffer)

    def on_data(self, s, data):
        self.request.write(data)

    def on_close(self, s):
        self.request.end()

    def write(self, data):
        if self.is_connected:
            self.data_len += len(data)
            self.conn.write(data)
        else:
            self.buffer = data

    def end(self):
        self.conn.end()

class Request(object):
    _requests = []

    def __init__(self, conn):
        self.conn = conn
        self.response = None
        self.protocol = None
        self.protocol_parse_end = False
        self.time = time.time()
        self.data_len = 0
        self.rbuffer, self.wbuffer = self.conn.buffer

        conn.on('data', self.on_data)
        conn.on('close', self.on_close)

    def parse(self, data, buffer):
        try:
            self.protocol.parse(data)
        except ProtocolParseEndError as e:
            self.protocol_parse_end = True
            if self.protocol.remote_addr.strip() == '0.0.0.0' and not self.protocol.remote_port:
                raise Exception("adder is empty %:%", self.protocol.remote_addr, self.protocol.remote_port)

            self.response = Response(self)
            if e.data:
                buffer.write(e.data)
                self.response.write(buffer)
            logging.info('%s connecting %s:%s %s', self.protocol, self.protocol.remote_addr, self.protocol.remote_port,
                         len(self._requests))
        except Exception:
            logging.error(traceback.format_exc())
            self.end()

    def on_data(self, s, data):
        if self.protocol_parse_end:
            return self.response.write(data)

        if self.protocol is None:
            self.protocol = HttpProtocol(self)
            self.parse(data.read(-1), data)
        else:
            self.parse(data.read(-1), data)

    def on_close(self, s):
        if self.response:
            self.response.end()
        self._requests.remove(self)
        logging.info('%s connected %s:%s %s %.3fs %s %s', self.protocol, self.protocol.remote_addr if self.protocol else '',
                     self.protocol.remote_port if self.protocol else '',
                     len(self._requests),
                     time.time()-self.time,
                     format_data_len(self.response.data_len if self.response else 0), format_data_len(self.data_len))
        self.response = None
        self.protocol = None

    def write(self, data):
        self.data_len += len(data)
        self.conn.write(data)

    def end(self):
        self.conn.end()

    @classmethod
    def on_connection(cls, s, conn):
        Request._requests.append(cls(conn))


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)1.1s %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S', filemode='a+')
    logging.info("starting server at port %s ..." % sys.argv[1])
    server = sevent.tcp.Server()
    server.enable_reuseaddr()
    server.listen(("0.0.0.0", int(sys.argv[1])))
    server.on('connection', Request.on_connection)
    sevent.current().start()