mirror of
https://github.com/zhigang1992/mitmproxy.git
synced 2026-01-12 17:32:27 +08:00
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:
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
"""
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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"])
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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)])
|
||||
|
||||
|
||||
Reference in New Issue
Block a user