From ec00b5a66e44454fce3d347203ad85ce4d8feaf6 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sun, 6 Mar 2011 16:54:49 +1300 Subject: [PATCH] Make mitmdump exit after client replay is complete by default. Add an option --keepserving to make it keep serving after replay. --- libmproxy/controller.py | 3 ++- libmproxy/dump.py | 6 +++++- libmproxy/flow.py | 20 ++++++++++++++++---- mitmdump | 16 ++++++++-------- test/test_flow.py | 14 ++++++++++---- 5 files changed, 41 insertions(+), 18 deletions(-) diff --git a/libmproxy/controller.py b/libmproxy/controller.py index 468092b74..82e934c4d 100644 --- a/libmproxy/controller.py +++ b/libmproxy/controller.py @@ -88,6 +88,7 @@ class Master: def shutdown(self): if not self._shutdown: self._shutdown = True - self.server.shutdown() + if self.server: + self.server.shutdown() diff --git a/libmproxy/dump.py b/libmproxy/dump.py index 621ffae57..b02c84dc7 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -15,6 +15,7 @@ class Options(object): "wfile", "rheaders", "stickycookie", + "keepserving", ] def __init__(self, **kwargs): for k, v in kwargs.items(): @@ -76,7 +77,10 @@ class DumpMaster(flow.FlowMaster): ) if options.client_replay: - self.start_client_playback(self._readflow(options.client_replay)) + self.start_client_playback( + self._readflow(options.client_replay), + not options.keepserving + ) def _readflow(self, path): path = os.path.expanduser(path) diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 857c8ae49..fb9bd4c72 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -32,13 +32,18 @@ class RequestReplayThread(threading.Thread): class ClientPlaybackState: - def __init__(self, flows): - self.flows = flows + def __init__(self, flows, exit): + self.flows, self.exit = flows, exit self.current = None def count(self): return len(self.flows) + def done(self): + if len(self.flows) == 0 and not self.current: + return True + return False + def clear(self, flow): """ A request has returned in some way - if this is the one we're @@ -447,11 +452,11 @@ class FlowMaster(controller.Master): else: self.stickycookie_state = None - def start_client_playback(self, flows): + def start_client_playback(self, flows, exit): """ flows: A list of flows. """ - self.client_playback = ClientPlaybackState(flows) + self.client_playback = ClientPlaybackState(flows, exit) def start_server_playback(self, flows, kill, headers): """ @@ -479,6 +484,13 @@ class FlowMaster(controller.Master): def tick(self, q): if self.client_playback: + e = [ + self.client_playback.done(), + self.client_playback.exit, + self.state.active_flow_count() == 0 + ] + if all(e): + self.shutdown() self.client_playback.tick(self) controller.Master.tick(self, q) diff --git a/mitmdump b/mitmdump index 3a44c6cd8..a8641d520 100755 --- a/mitmdump +++ b/mitmdump @@ -28,17 +28,19 @@ if __name__ == '__main__': usage = "%prog [options] [filter]", version="%%prog %s"%VERSION, ) - parser.add_option( - "-p", action="store", - type = "int", dest="port", default=8080, - help = "Proxy service port." - ) parser.add_option("-i", action="store_true", dest="stickycookie_all", default=None, help="Set sticky cookie for all requests.") parser.add_option("-I", action="store", dest="stickycookie_filt", default=None, metavar="FILTER", help="Set sticky cookie filter. Matched against requests.") + parser.add_option("--keepserving", + action="store_true", dest="keepserving", default=False, + help="Continue serving after playback.") + parser.add_option("-p", + action="store", + type = "int", dest="port", default=8080, + help = "Proxy service port.") parser.add_option("-q", action="store_true", dest="quiet", help="Quiet.") @@ -97,6 +99,7 @@ if __name__ == '__main__': kill = options.kill, rheaders = options.rheaders, client_replay = options.client_replay, + keepserving = options.keepserving, stickycookie = stickycookie, ) if args: @@ -104,9 +107,6 @@ if __name__ == '__main__': else: filt = None - if options.verbose > 0: - print >> sys.stderr, "Running on port %s"%options.port - try: m = dump.DumpMaster(server, dumpopts, filt) m.run() diff --git a/test/test_flow.py b/test/test_flow.py index b49d9cdf9..79e6dcc59 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -37,11 +37,12 @@ class uStickyCookieState(libpry.AutoTree): class uClientPlaybackState(libpry.AutoTree): def test_tick(self): first = tutils.tflow() - c = flow.ClientPlaybackState( - [first, tutils.tflow()] - ) s = flow.State() fm = flow.FlowMaster(None, s) + fm.start_client_playback([first, tutils.tflow()], True) + c = fm.client_playback + + assert not c.done() assert not s.flow_count() assert c.count() == 2 c.tick(fm, testing=True) @@ -54,7 +55,12 @@ class uClientPlaybackState(libpry.AutoTree): c.clear(c.current) c.tick(fm, testing=True) assert c.count() == 0 + c.clear(c.current) + assert c.done() + q = Queue.Queue() + fm.state.clear() + fm.tick(q) class uServerPlaybackState(libpry.AutoTree): @@ -391,7 +397,7 @@ class uFlowMaster(libpry.AutoTree): pb = [tutils.tflow_full(), f] fm = flow.FlowMaster(None, s) assert not fm.start_server_playback(pb, False, []) - assert not fm.start_client_playback(pb) + assert not fm.start_client_playback(pb, False) q = Queue.Queue() assert not fm.state.flow_count()