Make mitmdump exit after client replay is complete by default.

Add an option --keepserving to make it keep serving after replay.

libmproxy/controller.py

@@ -88,6 +88,7 @@ class Master:
     def shutdown(self):
         if not self._shutdown:
             self._shutdown = True
+            if self.server:
                 self.server.shutdown()
 
 

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)

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)
 

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()

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()