Rework client and server replay

- Add client.replay [flows], client.replay.stop
- Add server.replay [flows], server.replay.stop
- The corresponding options for file loading are only read on startup, further
changes are ignored. In interactive contexts, replay is started with the
commands, not through option changes.
- Deprecate flow.replay, use replay.client instead
This commit is contained in:
Aldo Cortesi
2017-04-29 11:02:36 +12:00
parent 7317ea134e
commit a92017a6c1
8 changed files with 71 additions and 81 deletions

View File

@@ -2,6 +2,7 @@ from mitmproxy import exceptions
from mitmproxy import ctx
from mitmproxy import io
from mitmproxy import flow
from mitmproxy import command
import typing
@@ -11,32 +12,46 @@ class ClientPlayback:
self.flows = None
self.current_thread = None
self.has_replayed = False
self.configured = False
def count(self) -> int:
if self.flows:
return len(self.flows)
return 0
def load(self, flows: typing.Sequence[flow.Flow]):
@command.command("replay.client.stop")
def stop_replay(self) -> None:
"""
Stop client replay.
"""
self.flows = []
ctx.master.addons.trigger("update", [])
@command.command("replay.client")
def start_replay(self, flows: typing.Sequence[flow.Flow]) -> None:
"""
Replay requests from flows.
"""
self.flows = flows
ctx.master.addons.trigger("update", [])
def configure(self, updated):
if "client_replay" in updated:
if ctx.options.client_replay:
ctx.log.info("Client Replay: {}".format(ctx.options.client_replay))
try:
flows = io.read_flows_from_paths(ctx.options.client_replay)
except exceptions.FlowReadException as e:
raise exceptions.OptionsError(str(e))
self.load(flows)
else:
self.flows = None
if not self.configured and ctx.options.client_replay:
self.configured = True
ctx.log.info("Client Replay: {}".format(ctx.options.client_replay))
try:
flows = io.read_flows_from_paths(ctx.options.client_replay)
except exceptions.FlowReadException as e:
raise exceptions.OptionsError(str(e))
self.start_replay(flows)
def tick(self):
if self.current_thread and not self.current_thread.is_alive():
self.current_thread = None
if self.flows and not self.current_thread:
self.current_thread = ctx.master.replay_request(self.flows.pop(0))
f = self.flows.pop(0)
self.current_thread = ctx.master.replay_request(f)
ctx.master.addons.trigger("update", [f])
self.has_replayed = True
if self.has_replayed:
if not self.flows and not self.current_thread:

View File

@@ -47,23 +47,12 @@ class Core:
@command.command("flow.mark.toggle")
def mark_toggle(self, flows: typing.Sequence[flow.Flow]) -> None:
"""
Mark flows.
Toggle mark for flows.
"""
for i in flows:
i.marked = not i.marked
ctx.master.addons.trigger("update", flows)
@command.command("flow.replay")
def replay(self, f: flow.Flow) -> None:
"""
Replay an HTTP flow request.
"""
try:
ctx.master.replay_request(f) # type: ignore
except exceptions.ReplayException as e:
raise exceptions.CommandError("Replay error: %s" % e) from e
ctx.master.addons.trigger("update", [f])
@command.command("flow.kill")
def kill(self, flows: typing.Sequence[flow.Flow]) -> None:
"""

View File

@@ -1,11 +1,14 @@
import hashlib
import urllib
import typing
from typing import Any # noqa
from typing import List # noqa
from mitmproxy import ctx
from mitmproxy import flow
from mitmproxy import exceptions
from mitmproxy import io
from mitmproxy import command
class ServerPlayback:
@@ -13,15 +16,27 @@ class ServerPlayback:
self.flowmap = {}
self.stop = False
self.final_flow = None
self.configured = False
def load_flows(self, flows):
@command.command("replay.server")
def load_flows(self, flows: typing.Sequence[flow.Flow]) -> None:
"""
Replay server responses from flows.
"""
self.flowmap = {}
for i in flows:
if i.response:
if i.response: # type: ignore
l = self.flowmap.setdefault(self._hash(i), [])
l.append(i)
ctx.master.addons.trigger("update", [])
def clear(self):
@command.command("replay.server.stop")
def clear(self) -> None:
"""
Stop server replay.
"""
self.flowmap = {}
ctx.master.addons.trigger("update", [])
def count(self):
return sum([len(i) for i in self.flowmap.values()])
@@ -90,14 +105,13 @@ class ServerPlayback:
return ret
def configure(self, updated):
if "server_replay" in updated:
self.clear()
if ctx.options.server_replay:
try:
flows = io.read_flows_from_paths(ctx.options.server_replay)
except exceptions.FlowReadException as e:
raise exceptions.OptionsError(str(e))
self.load_flows(flows)
if not self.configured and ctx.options.server_replay:
self.configured = True
try:
flows = io.read_flows_from_paths(ctx.options.server_replay)
except exceptions.FlowReadException as e:
raise exceptions.OptionsError(str(e))
self.load_flows(flows)
def tick(self):
if self.stop and not self.final_flow.live:

View File

@@ -132,14 +132,6 @@ class FlowItem(urwid.WidgetWrap):
def selectable(self):
return True
def server_replay_prompt(self, k):
a = self.master.addons.get("serverplayback")
if k == "a":
a.load([i.copy() for i in self.master.view])
elif k == "t":
a.load([self.flow.copy()])
signals.update_settings.send(self)
def mouse_event(self, size, event, button, col, row, focus):
if event == "mouse press" and button == 1:
if self.flow.request:
@@ -149,30 +141,7 @@ class FlowItem(urwid.WidgetWrap):
def keypress(self, xxx_todo_changeme, key):
(maxcol,) = xxx_todo_changeme
key = common.shortcuts(key)
if key == "S":
def stop_server_playback(response):
if response == "y":
self.master.options.server_replay = []
a = self.master.addons.get("serverplayback")
if a.count():
signals.status_prompt_onekey.send(
prompt = "Stop current server replay?",
keys = (
("yes", "y"),
("no", "n"),
),
callback = stop_server_playback,
)
else:
signals.status_prompt_onekey.send(
prompt = "Server Replay",
keys = (
("all flows", "a"),
("this flow", "t"),
),
callback = self.server_replay_prompt,
)
elif key == "V":
if key == "V":
if not self.flow.modified():
signals.status_message.send(message="Flow not modified.")
return

View File

@@ -78,7 +78,7 @@ class UnsupportedLog:
signals.add_log(strutils.bytes_to_escaped_str(message.content), "debug")
class ConsoleCommands:
class ConsoleAddon:
"""
An addon that exposes console-specific commands.
"""
@@ -131,6 +131,10 @@ class ConsoleCommands:
def running(self):
self.started = True
def update(self, flows):
if not flows:
signals.update_settings.send(self)
def configure(self, updated):
if self.started:
if "console_eventlog" in updated:
@@ -157,7 +161,8 @@ def default_keymap(km):
km.add("g", "view.go 0", context="flowlist")
km.add("G", "view.go -1", context="flowlist")
km.add("m", "flow.mark.toggle @focus", context="flowlist")
km.add("r", "flow.replay @focus", context="flowlist")
km.add("r", "replay.client @focus", context="flowlist")
km.add("S", "console.command 'replay.server '")
km.add("v", "set console_order_reversed=toggle", context="flowlist")
km.add("U", "flow.mark @all false", context="flowlist")
km.add("w", "console.command 'save.file @shown '", context="flowlist")
@@ -196,7 +201,7 @@ class ConsoleMaster(master.Master):
self.view,
UnsupportedLog(),
readfile.ReadFile(),
ConsoleCommands(self),
ConsoleAddon(self),
)
def sigint_handler(*args, **kwargs):

View File

@@ -26,7 +26,7 @@ class TestClientPlayback:
with taddons.context() as tctx:
assert cp.count() == 0
f = tflow.tflow(resp=True)
cp.load([f])
cp.start_replay([f])
assert cp.count() == 1
RP = "mitmproxy.proxy.protocol.http_replay.RequestReplayThread"
with mock.patch(RP) as rp:
@@ -44,13 +44,20 @@ class TestClientPlayback:
cp.tick()
assert cp.current_thread is None
cp.start_replay([f])
cp.stop_replay()
assert not cp.flows
def test_configure(self, tmpdir):
cp = clientplayback.ClientPlayback()
with taddons.context() as tctx:
path = str(tmpdir.join("flows"))
tdump(path, [tflow.tflow()])
tctx.configure(cp, client_replay=[path])
cp.configured = False
tctx.configure(cp, client_replay=[])
cp.configured = False
tctx.configure(cp)
cp.configured = False
with pytest.raises(exceptions.OptionsError):
tctx.configure(cp, client_replay=["nonexistent"])

View File

@@ -3,7 +3,6 @@ from mitmproxy.test import taddons
from mitmproxy.test import tflow
from mitmproxy import exceptions
import pytest
from unittest import mock
def test_set():
@@ -43,15 +42,6 @@ def test_mark():
assert f.marked
def test_replay():
sa = core.Core()
with taddons.context():
f = tflow.tflow()
with mock.patch("mitmproxy.master.Master.replay_request") as rp:
sa.replay(f)
assert rp.called
def test_kill():
sa = core.Core()
with taddons.context():

View File

@@ -22,6 +22,7 @@ def test_config(tmpdir):
fpath = str(tmpdir.join("flows"))
tdump(fpath, [tflow.tflow(resp=True)])
tctx.configure(s, server_replay=[fpath])
s.configured = False
with pytest.raises(exceptions.OptionsError):
tctx.configure(s, server_replay=[str(tmpdir)])