diff options
| -rw-r--r-- | libmproxy/controller.py | 3 | ||||
| -rw-r--r-- | libmproxy/dump.py | 6 | ||||
| -rw-r--r-- | libmproxy/flow.py | 20 | ||||
| -rwxr-xr-x | mitmdump | 16 | ||||
| -rw-r--r-- | test/test_flow.py | 14 |
5 files changed, 41 insertions, 18 deletions
diff --git a/libmproxy/controller.py b/libmproxy/controller.py index 468092b7..82e934c4 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 621ffae5..b02c84dc 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 857c8ae4..fb9bd4c7 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) @@ -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 b49d9cdf..79e6dcc5 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() |