From f908ea220f78754266a54dd57f34dcc24d761f1d Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@corte.si>
Date: Sat, 29 Apr 2017 11:34:50 +1200
Subject: [PATCH] commands: "replay.client.file" and "replay.server.file"

---
 mitmproxy/addons/clientplayback.py           |  8 ++++++++
 mitmproxy/addons/serverplayback.py           |  8 ++++++++
 test/mitmproxy/addons/test_clientplayback.py | 10 ++++++++++
 test/mitmproxy/addons/test_serverplayback.py | 11 +++++++++++
 4 files changed, 37 insertions(+)

diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py
index 322933f9c..0db6d336e 100644
--- a/mitmproxy/addons/clientplayback.py
+++ b/mitmproxy/addons/clientplayback.py
@@ -35,6 +35,14 @@ class ClientPlayback:
         self.flows = flows
         ctx.master.addons.trigger("update", [])
 
+    @command.command("replay.client.file")
+    def load_file(self, path: str) -> None:
+        try:
+            flows = io.read_flows_from_paths([path])
+        except exceptions.FlowReadException as e:
+            raise exceptions.CommandError(str(e))
+        self.flows = flows
+
     def configure(self, updated):
         if not self.configured and ctx.options.client_replay:
             self.configured = True
diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py
index ab318d4e8..927f6e15c 100644
--- a/mitmproxy/addons/serverplayback.py
+++ b/mitmproxy/addons/serverplayback.py
@@ -30,6 +30,14 @@ class ServerPlayback:
                 l.append(i)
         ctx.master.addons.trigger("update", [])
 
+    @command.command("replay.server.file")
+    def load_file(self, path: str) -> None:
+        try:
+            flows = io.read_flows_from_paths([path])
+        except exceptions.FlowReadException as e:
+            raise exceptions.CommandError(str(e))
+        self.load_flows(flows)
+
     @command.command("replay.server.stop")
     def clear(self) -> None:
         """
diff --git a/test/mitmproxy/addons/test_clientplayback.py b/test/mitmproxy/addons/test_clientplayback.py
index 843d7409d..7ffda317b 100644
--- a/test/mitmproxy/addons/test_clientplayback.py
+++ b/test/mitmproxy/addons/test_clientplayback.py
@@ -48,6 +48,16 @@ class TestClientPlayback:
             cp.stop_replay()
             assert not cp.flows
 
+    def test_load_file(self, tmpdir):
+        cp = clientplayback.ClientPlayback()
+        with taddons.context():
+            fpath = str(tmpdir.join("flows"))
+            tdump(fpath, [tflow.tflow(resp=True)])
+            cp.load_file(fpath)
+            assert cp.flows
+            with pytest.raises(exceptions.CommandError):
+                cp.load_file("/nonexistent")
+
     def test_configure(self, tmpdir):
         cp = clientplayback.ClientPlayback()
         with taddons.context() as tctx:
diff --git a/test/mitmproxy/addons/test_serverplayback.py b/test/mitmproxy/addons/test_serverplayback.py
index e0c025fe8..3ceab3fab 100644
--- a/test/mitmproxy/addons/test_serverplayback.py
+++ b/test/mitmproxy/addons/test_serverplayback.py
@@ -16,6 +16,17 @@ def tdump(path, flows):
         w.add(i)
 
 
+def test_load_file(tmpdir):
+    s = serverplayback.ServerPlayback()
+    with taddons.context():
+        fpath = str(tmpdir.join("flows"))
+        tdump(fpath, [tflow.tflow(resp=True)])
+        s.load_file(fpath)
+        assert s.flowmap
+        with pytest.raises(exceptions.CommandError):
+            s.load_file("/nonexistent")
+
+
 def test_config(tmpdir):
     s = serverplayback.ServerPlayback()
     with taddons.context() as tctx: