PyCodeBattler β

■ 名前

クローディア

■ ステータス

■ 必殺技

■ コード

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# $Id: proxy.py,v 1.6 2011/06/04 08:48:27 atzm Exp $

import sys
import string
import mimetools
import socket
import urllib
import SocketServer
import BaseHTTPServer


def spliturl(url):
    host, path = urllib.splithost(urllib.splittype(url)[1])
    host, port = urllib.splitport(host)
    if not port:
        port = 80
    port = int(port)
    return (host, port, path)


def write_header(hdr, wfile=sys.stdout, delim='\r\n'):
    for k, v in hdr.items():
        line = '%s: %s' % (string.capwords(k, '-'), v)
        wfile.write(line + delim)
    wfile.write(delim)


def handle_body(hdr, rfile=sys.stdin, wfile=sys.stdout, res=False):
    # XXX: 場合によっては刺さるかも :-P

    if hdr.get('transfer-encoding', '').lower() == 'chunked':
        while True:
            size_orig = rfile.readline()

            if not size_orig:
                break

            size = int(size_orig.strip(), 16)

            if size:
                wfile.write(size_orig)
                wfile.write(rfile.read(size + 2))  # + \r\n
            else:
                wfile.write(size_orig)
                break

    elif 'content-length' in hdr:
        bytes = hdr['content-length']
        data = rfile.read(int(bytes))
        wfile.write(data)

    elif res:
        for line in rfile:
            if not line:
                break
            wfile.write(line)


class ProxyHandler(BaseHTTPServer.BaseHTTPRequestHandler):
    def _proxy(self):
        host, port, path = spliturl(self.path)
        reqline = ' '.join([self.command, path, self.request_version])
        sock = socket.socket()

        try:
            sock.connect((host, port))
        except socket.error, e:
            print('[ERR] %s: %s:%d' % (e, host, port))
            return

        print('[DBG] connected to %s:%d' % (host, port))
        print('[DBG]   %s' % reqline)

        sfile = sock.makefile()
        sfile.write(reqline + '\r\n')
        write_header(self.headers, sfile)
        handle_body(self.headers, self.rfile, sfile)
        sfile.flush()

        resline = sfile.readline()
        reshead = mimetools.Message(sfile, 0)

        print('[DBG]   %s' % resline)

        self.wfile.write(resline)
        write_header(reshead, self.wfile)
        handle_body(reshead, sfile, self.wfile, True)
        self.wfile.flush()

        sfile.close()
        sock.close()

    def do_HEAD(self):
        self._proxy()

    def do_GET(self):
        self._proxy()

    def do_POST(self):
        self._proxy()


class ProxyServer(SocketServer.ForkingMixIn, BaseHTTPServer.HTTPServer):
    allow_reuse_address = True


def main():
    addr = sys.argv[1]
    port = int(sys.argv[2])
    server = ProxyServer((addr, port), ProxyHandler)
    server.serve_forever()


if __name__ == '__main__':
    main()