Add a way for users to specify header significance in server replay.

Also add the --rheader command-line option to mitmdump to let the user specify
an arbitrary number of significant headers. The default is to treat no headers
as significant.
This commit is contained in:
Aldo Cortesi 2011-02-23 10:54:51 +13:00
parent c80214ba55
commit 39207ffdd2
4 changed files with 61 additions and 20 deletions

View File

@ -12,6 +12,7 @@ class Options(object):
"replay",
"verbosity",
"wfile",
"rheaders",
]
def __init__(self, **kwargs):
for k, v in kwargs.items():
@ -52,7 +53,7 @@ class DumpMaster(flow.FlowMaster):
flows = list(flow.FlowReader(f).stream())
except IOError, v:
raise DumpError(v.strerror)
self.start_playback(flows, options.kill)
self.start_playback(flows, options.kill, options.rheaders)
def _runscript(self, f, script):
try:

View File

@ -32,7 +32,12 @@ class RequestReplayThread(threading.Thread):
class ServerPlaybackState:
def __init__(self):
def __init__(self, headers):
"""
headers: A case-insensitive list of request headers that should be
included in request-response matching.
"""
self.headers = headers
self.fmap = {}
def count(self):
@ -62,6 +67,15 @@ class ServerPlaybackState:
str(r.path),
str(r.content),
]
if self.headers:
hdrs = []
for i in self.headers:
v = r.headers.get(i, [])
# Slightly subtle: we need to convert everything to strings
# to prevent a mismatch between unicode/non-unicode.
v = [str(x) for x in v]
hdrs.append((i, v))
key.append(repr(hdrs))
return hashlib.sha256(repr(key)).digest()
def next_flow(self, request):
@ -342,12 +356,12 @@ class FlowMaster(controller.Master):
def set_request_script(self, s):
self.scripts["request"] = s
def start_playback(self, flows, kill):
def start_playback(self, flows, kill, headers):
"""
flows: A list of flows.
kill: Boolean, should we kill requests not part of the replay?
"""
self.playback = ServerPlaybackState()
self.playback = ServerPlaybackState(headers)
self.playback.load(flows)
self.kill_nonreplay = kill

View File

@ -28,31 +28,40 @@ if __name__ == '__main__':
)
proxy.certificate_option_group(parser)
parser.add_option(
"-p", "--port", action="store",
"-p", action="store",
type = "int", dest="port", default=8080,
help = "Port."
)
parser.add_option("-q", "--quiet",
parser.add_option("-q",
action="store_true", dest="quiet",
help="Quiet.")
parser.add_option("-v", "--verbose",
action="count", dest="verbose", default=1,
help="Increase verbosity. Can be passed multiple times.")
parser.add_option("-w", "--writefile",
action="store", dest="wfile", default=None,
help="Write flows to file.")
parser.add_option("", "--reqscript",
action="store", dest="request_script", default=None,
help="Script to run when a request is recieved.")
parser.add_option("", "--respscript",
action="store", dest="response_script", default=None,
help="Script to run when a response is recieved.")
parser.add_option("-r", "--replay",
action="store", dest="replay", default=None,
parser.add_option("-v",
action="count", dest="verbose", default=1,
help="Increase verbosity. Can be passed multiple times.")
parser.add_option("-w",
action="store", dest="wfile", default=None,
help="Write flows to file.")
group = OptionGroup(parser, "Server Replay")
group.add_option("-r", action="store", dest="replay", default=None, metavar="PATH",
help="Replay server responses from a saved file.")
parser.add_option("-k", "--kill",
group.add_option("-k", "--kill",
action="store_true", dest="kill", default=False,
help="Kill extra requests during replay.")
group.add_option("--rheader",
action="append", dest="rheaders", type="str",
help="Request headers to be considered during replay. "
"Can be passed multiple times.")
parser.add_option_group(group)
options, args = parser.parse_args()
@ -60,6 +69,7 @@ if __name__ == '__main__':
if options.quiet:
options.verbose = 0
config = proxy.process_certificate_option_group(parser, options)
server = proxy.ProxyServer(config, options.port)
dumpopts = dump.Options(
@ -68,7 +78,8 @@ if __name__ == '__main__':
request_script = options.request_script,
response_script = options.response_script,
replay = options.replay,
kill = options.kill
kill = options.kill,
rheaders = options.rheaders
)
if args:
filt = " ".join(args)

View File

@ -6,7 +6,7 @@ import libpry
class uServerPlaybackState(libpry.AutoTree):
def test_hash(self):
s = flow.ServerPlaybackState()
s = flow.ServerPlaybackState(None)
r = utils.tflow()
r2 = utils.tflow()
@ -17,8 +17,23 @@ class uServerPlaybackState(libpry.AutoTree):
r.request.path = "voing"
assert s._hash(r) != s._hash(r2)
def test_headers(self):
s = flow.ServerPlaybackState(["foo"])
r = utils.tflow_full()
r.request.headers["foo"] = ["bar"]
r2 = utils.tflow_full()
assert not s._hash(r) == s._hash(r2)
r2.request.headers["foo"] = ["bar"]
assert s._hash(r) == s._hash(r2)
r2.request.headers["oink"] = ["bar"]
assert s._hash(r) == s._hash(r2)
r = utils.tflow_full()
r2 = utils.tflow_full()
assert s._hash(r) == s._hash(r2)
def test_load(self):
s = flow.ServerPlaybackState()
s = flow.ServerPlaybackState(None)
r = utils.tflow_full()
r.request.headers["key"] = ["one"]
@ -319,10 +334,10 @@ class uFlowMaster(libpry.AutoTree):
fm = flow.FlowMaster(None, s)
assert not fm.do_playback(utils.tflow())
fm.start_playback(pb, False)
fm.start_playback(pb, False, [])
assert fm.do_playback(utils.tflow())
fm.start_playback(pb, False)
fm.start_playback(pb, False, [])
r = utils.tflow()
r.request.content = "gibble"
assert not fm.do_playback(r)