diff options
Diffstat (limited to 'libmproxy')
| -rw-r--r-- | libmproxy/dump.py | 3 | ||||
| -rw-r--r-- | libmproxy/flow.py | 20 |
2 files changed, 19 insertions, 4 deletions
diff --git a/libmproxy/dump.py b/libmproxy/dump.py index 66bb5206..396dc31d 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -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: diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 80a88708..42870f17 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -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 |