Make mitmdump exit after client replay is complete by default.

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

3 files changed, 23 insertions, 6 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)