Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 00:01:36 +00:00
Code Issues Projects Releases Wiki Activity

Commands, core update event

Browse Source 
This patch:

- Introduces a core update() event that should be invoked whenever flows are
changed outside of the normal lifecycle.
- Extend view.resolve to know about @all, which matches all flows in the view.
- Add a core flow.resume comand, which resumes flows and broadcasts an update event.
- Define flow list bindings for:

A -> flow.resume @all
a -> flow.resume @focus
d -> view.remove @focus
z -> view.remove @all
This commit is contained in:
Aldo Cortesi 2017-04-28 17:01:48 +12:00
parent b537997f4f
commit 0b090f7ae1
10 changed files with 104 additions and 89 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
mitmproxy/addons/core.py
View File

@ -1,6 +1,9 @@
import typing
from mitmproxy import ctx
from mitmproxy import exceptions
from mitmproxy import command
from mitmproxy import flow
class Core:
@ -16,3 +19,13 @@ class Core:
ctx.options.set(spec)
except exceptions.OptionsError as e:
raise exceptions.CommandError(e) from e
@command.command("flow.resume")
def resume(self, flows: typing.Sequence[flow.Flow]) -> None:
"""
Resume flows if they are intercepted.
"""
intercepted = [i for i in flows if i.intercepted]
for f in intercepted:
f.resume()
ctx.master.addons.trigger("update", intercepted)

78
mitmproxy/addons/view.py
View File

@ -18,6 +18,7 @@ import sortedcontainers
import mitmproxy.flow
from mitmproxy import flowfilter
from mitmproxy import exceptions
from mitmproxy import command
from mitmproxy import ctx
from mitmproxy import http # noqa
@ -223,7 +224,7 @@ class View(collections.Sequence):
self.filter = flt or matchall
self._refilter()
def clear(self):
def clear(self) -> None:
"""
Clears both the store and view.
"""
@ -256,43 +257,21 @@ class View(collections.Sequence):
self.focus.flow = f
self.sig_view_add.send(self, flow=f)
def remove(self, f: mitmproxy.flow.Flow):
@command.command("view.remove")
def remove(self, flows: typing.Sequence[mitmproxy.flow.Flow]) -> None:
"""
Removes the flow from the underlying store and the view.
"""
for f in flows:
if f.id in self._store:
if f.killable:
f.kill()
if f in self._view:
self._view.remove(f)
self.sig_view_remove.send(self, flow=f)
del self._store[f.id]
self.sig_store_remove.send(self, flow=f)
def update(self, f: mitmproxy.flow.Flow):
"""
Updates a flow. If the flow is not in the state, it's ignored.
"""
if f.id in self._store:
if self.filter(f):
if f not in self._view:
self._base_add(f)
if self.focus_follow:
self.focus.flow = f
self.sig_view_add.send(self, flow=f)
else:
# This is a tad complicated. The sortedcontainers
# implementation assumes that the order key is stable. If
# it changes mid-way Very Bad Things happen. We detect when
# this happens, and re-fresh the item.
self.order_key.refresh(f)
self.sig_view_update.send(self, flow=f)
else:
try:
self._view.remove(f)
self.sig_view_remove.send(self, flow=f)
except ValueError:
# The value was not in the view
pass
def get_by_id(self, flow_id: str) -> typing.Optional[mitmproxy.flow.Flow]:
"""
Get flow with the given id from the store.
@ -322,10 +301,13 @@ class View(collections.Sequence):
if "console_focus_follow" in updated:
self.focus_follow = ctx.options.console_focus_follow
@command.command("view.resolve")
def resolve(self, spec: str) -> typing.Sequence[mitmproxy.flow.Flow]:
"""
Resolve a flow list specification to an actual list of flows.
"""
if spec == "@all":
return [i for i in self._store.values()]
if spec == "@focus":
return [self.focus.flow] if self.focus.flow else []
elif spec == "@shown":
@ -342,26 +324,50 @@ class View(collections.Sequence):
raise exceptions.CommandError("Invalid flow filter: %s" % spec)
return [i for i in self._store.values() if filt(i)]
def load(self, l):
l.add_command("console.resolve", self.resolve)
def request(self, f):
self.add(f)
def error(self, f):
self.update(f)
self.update([f])
def response(self, f):
self.update(f)
self.update([f])
def intercept(self, f):
self.update(f)
self.update([f])
def resume(self, f):
self.update(f)
self.update([f])
def kill(self, f):
self.update(f)
self.update([f])
def update(self, flows: typing.Sequence[mitmproxy.flow.Flow]) -> None:
"""
Updates a list of flows. If flow is not in the state, it's ignored.
"""
for f in flows:
if f.id in self._store:
if self.filter(f):
if f not in self._view:
self._base_add(f)
if self.focus_follow:
self.focus.flow = f
self.sig_view_add.send(self, flow=f)
else:
# This is a tad complicated. The sortedcontainers
# implementation assumes that the order key is stable. If
# it changes mid-way Very Bad Things happen. We detect when
# this happens, and re-fresh the item.
self.order_key.refresh(f)
self.sig_view_update.send(self, flow=f)
else:
try:
self._view.remove(f)
self.sig_view_remove.send(self, flow=f)
except ValueError:
# The value was not in the view
pass
class Focus:

4
mitmproxy/command.py
View File

@ -110,9 +110,9 @@ def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any:
if argtype == str:
return spec
elif argtype == typing.Sequence[flow.Flow]:
return manager.call_args("console.resolve", [spec])
return manager.call_args("view.resolve", [spec])
elif argtype == flow.Flow:
flows = manager.call_args("console.resolve", [spec])
flows = manager.call_args("view.resolve", [spec])
if len(flows) != 1:
raise exceptions.CommandError(
"Command requires one flow, specification matched %s." % len(flows)

1
mitmproxy/eventsequence.py
View File

@ -35,6 +35,7 @@ Events = frozenset([
"load",
"running",
"tick",
"update",
])

34
mitmproxy/tools/console/flowlist.py
View File

@ -150,14 +150,7 @@ class FlowItem(urwid.WidgetWrap):
def keypress(self, xxx_todo_changeme, key):
(maxcol,) = xxx_todo_changeme
key = common.shortcuts(key)
if key == "a":
self.flow.resume()
self.master.view.update(self.flow)
elif key == "d":
if self.flow.killable:
self.flow.kill()
self.master.view.remove(self.flow)
elif key == "D":
if key == "D":
cp = self.flow.copy()
self.master.view.add(cp)
self.master.view.focus.flow = cp
@ -222,14 +215,14 @@ class FlowItem(urwid.WidgetWrap):
callback = common.export_to_clip_or_file,
args = (None, self.flow, common.ask_save_path)
)
elif key == "C":
signals.status_prompt_onekey.send(
self,
prompt = "Export to clipboard",
keys = [(e[0], e[1]) for e in export.EXPORTERS],
callback = common.export_to_clip_or_file,
args = (None, self.flow, common.copy_to_clipboard_or_prompt)
)
# elif key == "C":
# signals.status_prompt_onekey.send(
# self,
# prompt = "Export to clipboard",
# keys = [(e[0], e[1]) for e in export.EXPORTERS],
# callback = common.export_to_clip_or_file,
# args = (None, self.flow, common.copy_to_clipboard_or_prompt)
# )
elif key == "b":
common.ask_save_body(None, self.flow)
else:
@ -321,14 +314,7 @@ class FlowListBox(urwid.ListBox):
def keypress(self, size, key):
key = common.shortcuts(key)
if key == "A":
for f in self.master.view:
if f.intercepted:
f.resume()
self.master.view.update(f)
elif key == "z":
self.master.view.clear()
elif key == "Z":
if key == "Z":
self.master.view.clear_not_marked()
elif key == "g":
if len(self.master.view):

4
mitmproxy/tools/console/master.py
View File

@ -147,11 +147,15 @@ def default_keymap(km):
km.add("i", "console.command 'set intercept='")
km.add("W", "console.command 'set save_stream_file='")
km.add("A", "flow.resume @all", context="flowlist")
km.add("a", "flow.resume @focus", context="flowlist")
km.add("d", "view.remove @focus", context="flowlist")
km.add("F", "set console_focus_follow=toggle", context="flowlist")
km.add("v", "set console_order_reversed=toggle", context="flowlist")
km.add("f", "console.command 'set view_filter='", context="flowlist")
km.add("e", "set console_eventlog=toggle", context="flowlist")
km.add("w", "console.command 'save.file @shown '", context="flowlist")
km.add("z", "view.remove @all", context="flowlist")
km.add("enter", "console.view.flow @focus", context="flowlist")

18
mitmproxy/tools/web/app.py
View File

@ -246,7 +246,7 @@ class ResumeFlows(RequestHandler):
def post(self):
for f in self.view:
f.resume()
self.view.update(f)
self.view.update([f])
class KillFlows(RequestHandler):
@ -254,27 +254,27 @@ class KillFlows(RequestHandler):
for f in self.view:
if f.killable:
f.kill()
self.view.update(f)
self.view.update([f])
class ResumeFlow(RequestHandler):
def post(self, flow_id):
self.flow.resume()
self.view.update(self.flow)
self.view.update([self.flow])
class KillFlow(RequestHandler):
def post(self, flow_id):
if self.flow.killable:
self.flow.kill()
self.view.update(self.flow)
self.view.update([self.flow])
class FlowHandler(RequestHandler):
def delete(self, flow_id):
if self.flow.killable:
self.flow.kill()
self.view.remove(self.flow)
self.view.remove([self.flow])
def put(self, flow_id):
flow = self.flow
@ -317,7 +317,7 @@ class FlowHandler(RequestHandler):
except APIError:
flow.revert()
raise
self.view.update(flow)
self.view.update([flow])
class DuplicateFlow(RequestHandler):
@ -331,14 +331,14 @@ class RevertFlow(RequestHandler):
def post(self, flow_id):
if self.flow.modified():
self.flow.revert()
self.view.update(self.flow)
self.view.update([self.flow])
class ReplayFlow(RequestHandler):
def post(self, flow_id):
self.flow.backup()
self.flow.response = None
self.view.update(self.flow)
self.view.update([self.flow])
try:
self.master.replay_request(self.flow)
@ -351,7 +351,7 @@ class FlowContent(RequestHandler):
self.flow.backup()
message = getattr(self.flow, message)
message.content = self.filecontents
self.view.update(self.flow)
self.view.update([self.flow])
def get(self, flow_id, message):
message = getattr(self.flow, message)

29
test/mitmproxy/addons/test_view.py
View File

@ -31,7 +31,7 @@ def test_order_refresh():
v.add(tf)
tf.request.timestamp_start = 1
assert not sargs
v.update(tf)
v.update([tf])
assert sargs
@ -140,6 +140,7 @@ def test_load():
def test_resolve():
v = view.View()
with taddons.context(options=options.Options()) as tctx:
assert tctx.command(v.resolve, "@all") == []
assert tctx.command(v.resolve, "@focus") == []
assert tctx.command(v.resolve, "@shown") == []
assert tctx.command(v.resolve, "@hidden") == []
@ -149,6 +150,7 @@ def test_resolve():
v.request(tft(method="get"))
assert len(tctx.command(v.resolve, "~m get")) == 1
assert len(tctx.command(v.resolve, "@focus")) == 1
assert len(tctx.command(v.resolve, "@all")) == 1
assert len(tctx.command(v.resolve, "@shown")) == 1
assert len(tctx.command(v.resolve, "@unmarked")) == 1
assert tctx.command(v.resolve, "@hidden") == []
@ -156,6 +158,7 @@ def test_resolve():
v.request(tft(method="put"))
assert len(tctx.command(v.resolve, "@focus")) == 1
assert len(tctx.command(v.resolve, "@shown")) == 2
assert len(tctx.command(v.resolve, "@all")) == 2
assert tctx.command(v.resolve, "@hidden") == []
assert tctx.command(v.resolve, "@marked") == []
@ -175,6 +178,7 @@ def test_resolve():
assert m(tctx.command(v.resolve, "@hidden")) == ["PUT", "PUT"]
assert m(tctx.command(v.resolve, "@marked")) == ["GET"]
assert m(tctx.command(v.resolve, "@unmarked")) == ["PUT", "GET", "PUT"]
assert m(tctx.command(v.resolve, "@all")) == ["GET", "PUT", "GET", "PUT"]
with pytest.raises(exceptions.CommandError, match="Invalid flow filter"):
tctx.command(v.resolve, "~")
@ -230,14 +234,14 @@ def test_update():
assert f in v
f.request.method = "put"
v.update(f)
v.update([f])
assert f not in v
f.request.method = "get"
v.update(f)
v.update([f])
assert f in v
v.update(f)
v.update([f])
assert f in v
@ -291,14 +295,14 @@ def test_signals():
# An update that results in a flow being added to the view
clearrec()
v[0].request.method = "PUT"
v.update(v[0])
v.update([v[0]])
assert rec_remove
assert not any([rec_update, rec_refresh, rec_add])
# An update that does not affect the view just sends update
v.set_filter(flowfilter.parse("~m put"))
clearrec()
v.update(v[0])
v.update([v[0]])
assert rec_update
assert not any([rec_remove, rec_refresh, rec_add])
@ -307,7 +311,7 @@ def test_signals():
v.set_filter(flowfilter.parse("~m get"))
assert not len(v)
clearrec()
v.update(f)
v.update([f])
assert not any([rec_add, rec_update, rec_remove, rec_refresh])
@ -333,7 +337,7 @@ def test_focus_follow():
assert v.focus.flow.request.timestamp_start == 7
mod.request.method = "GET"
v.update(mod)
v.update([mod])
assert v.focus.index == 2
assert v.focus.flow.request.timestamp_start == 6
@ -374,15 +378,16 @@ def test_focus():
assert f.index == 0
f.index = 1
v.remove(v[1])
v.remove([v[1]])
v[1].intercept()
assert f.index == 1
assert f.flow is v[1]
v.remove(v[1])
v.remove([v[1]])
assert f.index == 0
assert f.flow is v[0]
v.remove(v[0])
v.remove([v[0]])
assert f.index is None
assert f.flow is None
@ -413,7 +418,7 @@ def test_settings():
v.settings[f]["foo"] = "bar"
assert v.settings[f]["foo"] == "bar"
assert len(list(v.settings)) == 1
v.remove(f)
v.remove([f])
with pytest.raises(KeyError):
v.settings[f]
assert not v.settings.keys()

2
test/mitmproxy/test_command.py
View File

@ -65,7 +65,7 @@ def test_typename():
class DummyConsole:
def load(self, l):
l.add_command("console.resolve", self.resolve)
l.add_command("view.resolve", self.resolve)
def resolve(self, spec: str) -> typing.Sequence[flow.Flow]:
n = int(spec)

2
test/mitmproxy/tools/web/test_app.py
View File

@ -162,7 +162,7 @@ class TestApp(tornado.testing.AsyncHTTPTestCase):
f = self.view.get_by_id(resp.body.decode())
assert f
assert f.id != "42"
self.view.remove(f)
self.view.remove([f])
def test_flow_revert(self):
f = self.view.get_by_id("42")