diff options
author | Aldo Cortesi <aldo@nullcube.com> | 2011-02-21 11:08:35 +1300 |
---|---|---|
committer | Aldo Cortesi <aldo@nullcube.com> | 2011-02-21 11:08:35 +1300 |
commit | fe99871df873f755ef5f4770edf37304f88102d7 (patch) | |
tree | 81cac2d8573f4e48d08284baadff8e0440e58ac0 | |
parent | c3e38970718aed37dd70e8aad0085957b62a09ac (diff) | |
download | mitmproxy-fe99871df873f755ef5f4770edf37304f88102d7.tar.gz mitmproxy-fe99871df873f755ef5f4770edf37304f88102d7.tar.bz2 mitmproxy-fe99871df873f755ef5f4770edf37304f88102d7.zip |
Add --kill option to mitmdump
If this option is passed all requests that are not part of a replayed
conversation are killed. If the option is not passed, such requests are passed
through to the server as usual.
-rw-r--r-- | libmproxy/dump.py | 15 | ||||
-rw-r--r-- | libmproxy/flow.py | 7 | ||||
-rwxr-xr-x | mitmdump | 4 | ||||
-rw-r--r-- | test/test_dump.py | 14 | ||||
-rw-r--r-- | test/test_flow.py | 4 | ||||
-rw-r--r-- | test/utils.py | 7 |
6 files changed, 40 insertions, 11 deletions
diff --git a/libmproxy/dump.py b/libmproxy/dump.py index d43da44b..7f6f1e7c 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -6,11 +6,12 @@ class DumpError(Exception): pass class Options(object): __slots__ = [ - "verbosity", - "wfile", + "kill", "request_script", "response_script", "replay", + "verbosity", + "wfile", ] def __init__(self, **kwargs): for k, v in kwargs.items(): @@ -66,7 +67,15 @@ class DumpMaster(flow.FlowMaster): f = flow.FlowMaster.handle_request(self, r) if self.o.request_script: self._runscript(f, self.o.request_script) - if not self.playback(f): + + if self.o.replay: + pb = self.playback(f) + if not pb: + if self.o.kill: + self.state.kill_flow(f) + else: + r.ack() + else: r.ack() def indent(self, n, t): diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 4025a30d..996b45cd 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -44,9 +44,10 @@ class ServerPlaybackState: chronological order. """ for i in flows: - h = self._hash(i) - l = self.fmap.setdefault(self._hash(i), []) - l.append(i) + if i.response: + h = self._hash(i) + l = self.fmap.setdefault(self._hash(i), []) + l.append(i) def _hash(self, flow): """ @@ -50,6 +50,9 @@ if __name__ == '__main__': parser.add_option("-r", "--replay", action="store", dest="replay", default=None, help="Replay server responses from a saved file.") + parser.add_option("-k", "--kill", + action="store_true", dest="kill", default=False, + help="Kill extra requests during replay.") options, args = parser.parse_args() @@ -65,6 +68,7 @@ if __name__ == '__main__': request_script = options.request_script, response_script = options.response_script, replay = options.replay, + kill = options.kill ) if args: filt = " ".join(args) diff --git a/test/test_dump.py b/test/test_dump.py index 90540408..6213f870 100644 --- a/test/test_dump.py +++ b/test/test_dump.py @@ -26,19 +26,27 @@ class uDumpMaster(libpry.AutoTree): def test_replay(self): cs = StringIO() - o = dump.Options(replay="nonexistent") + o = dump.Options(replay="nonexistent", kill=True) libpry.raises(dump.DumpError, dump.DumpMaster, None, o, None, outfile=cs) t = self.tmpdir() p = os.path.join(t, "rep") f = open(p, "w") fw = flow.FlowWriter(f) - t = utils.tflow() + t = utils.tflow_full() + t.response = utils.tresp(t.request) fw.add(t) f.close() - o = dump.Options(replay=p) + o = dump.Options(replay=p, kill=True) m = dump.DumpMaster(None, o, None, outfile=cs) + + self._cycle(m, "content") + self._cycle(m, "content") + + o = dump.Options(replay=p, kill=False) + m = dump.DumpMaster(None, o, None, outfile=cs) + self._cycle(m, "nonexistent") def test_options(self): o = dump.Options(verbosity = 2) diff --git a/test/test_flow.py b/test/test_flow.py index cd88464d..4bfd5ad9 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -19,10 +19,10 @@ class uServerPlaybackState(libpry.AutoTree): def test_load(self): s = flow.ServerPlaybackState() - r = utils.tflow() + r = utils.tflow_full() r.request.headers["key"] = ["one"] - r2 = utils.tflow() + r2 = utils.tflow_full() r2.request.headers["key"] = ["two"] s.load([r, r2]) diff --git a/test/utils.py b/test/utils.py index 12646106..9a00983c 100644 --- a/test/utils.py +++ b/test/utils.py @@ -22,3 +22,10 @@ def tflow(): return flow.Flow(r) +def tflow_full(): + r = treq() + f = flow.Flow(r) + f.response = tresp(r) + return f + + |