methaneのブログ

このブログに乗せているサンプルコードはすべてNYSLです。

gevent で proxy server を作る

Flash製アプリの試験用のアクセスパターンを作りたくて gevent を使って簡単なproxy serverを作った。

localhostの8000版で立ち上がるので、FoxyProxy などで、特定のドメインに対する proxy として登録してやると、リクエストを再現するのに必要なデータをmsgpackフォーマットで保存していく。

from gevent.pywsgi import WSGIServer
from gevent import monkey
from msgpack import packs
import httplib

def handler(environ, start_response):
    host = environ['HTTP_HOST']
    method = environ['REQUEST_METHOD']
    body = environ['wsgi.input'].read()
    path = environ['PATH_INFO']
    qs = environ['QUERY_STRING']
    content_type = environ.get('CONTENT_TYPE')

    headers = {}
    if content_type:
        headers['Content-Type'] = content_type
    for k,v in environ.iteritems():
        if not k.startswith('HTTP_'):
            continue
        headers[k[5:]] = v

    con = httplib.HTTPConnection(host)
    con.request(
            method,
            path + '?' + qs,
            body,
            headers,
            )
    REQ_LOG.write(packs(
        (host, method, path+'?'+qs, body, headers)
        ))

    print (host, method, path+'?'+qs, body, headers)

    res = con.getresponse()

    headers = res.getheaders()


    headers = [(k,v) for (k,v) in headers
            if k not in ('transfer-encoding', 'content-length')]

    status = "%d %s" % (res.status, res.reason)
    start_response(status, headers)
    return [res.read()]

def main():
    global REQ_LOG
    REQ_LOG = open('request_log.mpac', 'wb')
    try:
        monkey.patch_socket()
        server = WSGIServer(('', 8000), handler)
        server.serve_forever()
    finally:
        REQ_LOG.close()


if __name__ == '__main__':
    main()