Duplicate/Replay and Reroute HTTP Request and Compare Responses with mitmproxy

replay_and_test.py

# File: Vagrantfile
#   Vagrant.configure(2) do |config|
#     config.vm.box = "ubuntu/trusty64"
#   end
#
# Process:
#   $ wget https://gist.githubusercontent.com/nextrevision/f81cc9493af565613e6a/raw/e48f4229743f1f29e82fb17d7129c8dead3c51f0/replay_and_test.py
#   $ vagrant up
#   $ vagrant ssh
#   $ sudo apt-get update
#   $ sudo apt-get -y install mitmproxy
#   $ python -m SimpleHTTPServer 8000 &
#   $ python -m SimpleHTTPServer 8001 &
#   $ mitmdump -s /vagrant/replay_and_test.py &
#   $ curl -x http://localhost:8080 http://localhost:8000 # success
#   $ kill %2
#   $ cd /tmp && python -m SimpleHTTPServer 8001 &
#   $ curl -x http://localhost:8080 http://localhost:8000 # fail
#
import re
import uuid

state = {}
matcher = re.compile('<li><a href=".bashrc">.bashrc</a>')

def _all_finished(tracker_state, items):
    for item in items:
        if False in [v[item] for k,v in tracker_state.iteritems()]:
            return False
    return True

def _compare(tracker_state):
    values = [v for v in tracker_state.values()]
    expected_values = len(values[0].keys())
    value_set = (set(v.items()) for v in values)
    value_set_r = reduce(set.intersection, value_set)
    total_values = [v[1] for v in value_set_r]
    if len(total_values) == expected_values:
        print "Success: %s" % str(values[0]['code'])
    else:
        print "Fail:"
        for host, result in tracker_state.iteritems():
            print "  %s: code=%s match=%s" % (host, str(result['code']), str(result['match']))

def request(context, flow):
    global state
    tracker_id = str(uuid.uuid4())
    state[tracker_id] = {
        'orig': {
            'code': False,
            'match': False
        },
        'dup': {
            'code': False,
            'match': False
        }
    }
    flow.request.headers["tracker"] = [tracker_id]
    flow.request.headers["host_id"] = ['orig']
    flow_dup = context.duplicate_flow(flow)
    flow_dup.request.port = 8001
    flow_dup.request.headers["host_id"] = ['dup']
    context.replay_request(flow_dup)

def response(context, flow):
    global state
    global matcher
    match = False
    code = flow.response.code
    host_id = flow.request.headers['host_id'][0]
    if matcher.search(flow.response.content):
        match = True
    if 'tracker' in flow.request.headers.keys():
        tracker_id = flow.request.headers['tracker'][0]
        if tracker_id in state.keys():
            state[tracker_id][host_id]['code'] = code
            state[tracker_id][host_id]['match'] = match
            if _all_finished(state[tracker_id], ['code']):
                _compare(state[tracker_id])
                del state[tracker_id]