#!/usr/bin/env python

import argparse
import BaseHTTPServer
import urllib
import urllib2

class NoRedirectHandler(urllib2.HTTPRedirectHandler):
    def http_error_302(self, req, fp, code, msg, headers):
        infourl = urllib.addinfourl(fp, headers, req.get_full_url())
        infourl.status = code
        infourl.code = code
        return infourl
    http_error_300 = http_error_302
    http_error_301 = http_error_302
    http_error_303 = http_error_302
    http_error_307 = http_error_302


opener = urllib2.build_opener(NoRedirectHandler())
urllib2.install_opener(opener)


class UserProxyHandler(BaseHTTPServer.BaseHTTPRequestHandler):

    @property
    def dest_url(self):
        return "http://localhost:%s%s" % (self.server.args.backend_port, self.path)

    def finish_request(self, code, headers, output):
        self.send_response(code)
        for key, value in headers:
            self.send_header(key, value)
        self.end_headers()
        self.wfile.writelines(output)

    def updated_headers(self):
        headers = dict(self.headers)
        for header in self.server.args.header:
            key, value = header.split(":", 1)
            headers[key.strip()] = value.strip()
        return headers

    def do_request(self, request):
        try:
            url = urllib2.urlopen(request)
            data = url.readlines()
        except urllib2.HTTPError as err:
            return self.finish_request(err.getcode(), err.info().items(), err.readlines())
        except urllib2.URLError as err:
            return self.finish_request(503, err.info().items(), "503 Service Unavailable: %s" % err)

        return self.finish_request(url.getcode(), url.info().items(), data)

    def do_GET(self):
        headers = self.updated_headers()
        request = urllib2.Request(self.dest_url, headers=headers)
        self.do_request(request)

    def do_POST(self):
        content_len = int(self.headers['Content-Length'])
        data = self.rfile.read(content_len)
        headers = self.updated_headers()
        request = urllib2.Request(self.dest_url, headers=headers, data=data)
        self.do_request(request)


def main():

    description_msg = "Mediocre Reverse Proxy."
    parser = argparse.ArgumentParser(description=description_msg)

    parser.add_argument("-p", "--listen-port", default=8888, type=int,
                        help="Port to listen on.")
    parser.add_argument("-P", "--backend-port", default=8989, type=int,
                        help="Port to proxy to.")
    parser.add_argument("--header", default=[], action="append",
                        help="Headers to pass to backend. e.g. --header")

    args = parser.parse_args()

    # localhost is hardcoded to discourage real use. :)
    server = BaseHTTPServer.HTTPServer(('localhost', args.listen_port), UserProxyHandler)
    server.args = args
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print "Bye!"


if __name__ == "__main__":
    main()
