Simple Flask Proxy

app.py

from flask import Flask, request, abort
import hashlib
import json
import requests
import os

app = Flask(__name__)

app.config['CAPTURE'] = bool(os.environ.get('CAPTURE'))
app.config['CACHE_DIR'] = os.environ.get('CACHE_DIR', '/tmp/cache')


@app.route("/")
def hello():
    url = request.args.get('url')
    if not url or url.startswith('http://localhost'):
        return abort(400)

    return _retrieve(url)


def _retrieve(url):
    path = _cache_path(url)

    if not os.path.exists(path):
        if app.config['CAPTURE']:
            response = _request(url)
            _store(url, response)
        else:
            return abort(404)

    with open(path, 'r') as f:
        cached = json.loads(f.read())
    return cached['body'], cached['status'], cached['headers']


def _request(url):
    kwargs = __adapt_request_args(url)
    r = requests.request(request.method, url, **kwargs)
    return r


def __adapt_request_args(url):
    kwargs = {
        'data': request.data,
        'allow_redirects': False
    }
    headers = dict([(key.upper(), value)
                    for key, value in request.headers.items()])

    # Explicitly set content-length request header
    if 'CONTENT-LENGTH' not in headers or not headers['CONTENT-LENGTH']:
        headers['CONTENT-LENGTH'] = str(len(kwargs['data']))

    # Let requests reset the host for us.
    if 'HOST' in headers:
        del headers['HOST']

    kwargs['headers'] = headers
    return kwargs


def _store(url, response):
    cached = json.dumps({
        'body': response.text,
        'headers': __process_response_headers(response),
        'status': response.status_code
    })

    path = _cache_path(url)
    dir_ = path.rsplit('/', 1)[0]
    if not os.path.isdir(dir_):
        os.makedirs(dir_)
    with open(path, 'w') as w:
        w.write(cached)


def __process_response_headers(response):
    headers = dict(response.headers)
    headers['content-type'] = 'text/html; charset=utf-8'
    if 'content-encoding' in headers:
        del headers['content-encoding']
    if 'transfer-encoding' in headers:
        del headers['transfer-encoding']
    return headers


def _cache_path(url):
    to_hash = "%s%s%s" % (request.method, url, request.data)
    hashed = hashlib.md5(to_hash).hexdigest()
    address = (hashed[:2], hashed[2:4], hashed)
    return '%s/%s/%s/%s.json' % (app.config['CACHE_DIR'], address)


if __name__ == "__main__":
    app.run(port=4777)