mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-26 02:10:59 +00:00
Add saving and loading of complete flows for later replay and analysis.
This commit is contained in:
Aldo Cortesi
parent
29d8007678
commit
635f7a971d
@ -73,8 +73,8 @@ class ConnectionItem(WWrap):
|
||||
self.intercepting = True
|
||||
self.w = self.get_text()
|
||||
|
||||
def get_text(self, nofocus=False):
|
||||
return urwid.Text(self.flow.get_text(nofocus))
|
||||
def get_text(self):
|
||||
return urwid.Text(self.flow.get_text())
|
||||
|
||||
def selectable(self):
|
||||
return True
|
||||
@ -278,23 +278,6 @@ class ConnectionView(WWrap):
|
||||
self.flow.request.method = i[0].upper()
|
||||
self.master.refresh_connection(self.flow)
|
||||
|
||||
def save_connection(self, path):
|
||||
if not path:
|
||||
return
|
||||
if self.viewing == self.REQ:
|
||||
c = self.flow.request
|
||||
else:
|
||||
c = self.flow.response
|
||||
path = os.path.expanduser(path)
|
||||
try:
|
||||
f = file(path, "w")
|
||||
f.write(str(c.headers))
|
||||
f.write("\r\n")
|
||||
f.write(str(c.content))
|
||||
f.close()
|
||||
except IOError, v:
|
||||
self.master.statusbar.message(str(v))
|
||||
|
||||
def edit(self, part):
|
||||
if self.viewing == self.REQ:
|
||||
conn = self.flow.request
|
||||
@ -368,10 +351,7 @@ class ConnectionView(WWrap):
|
||||
self.state.revert(self.flow)
|
||||
self.master.refresh_connection(self.flow)
|
||||
elif key == "S":
|
||||
if self.viewing == self.REQ:
|
||||
self.master.prompt("Save request: ", self.save_connection)
|
||||
else:
|
||||
self.master.prompt("Save response: ", self.save_connection)
|
||||
self.master.prompt("Save all: ", self.save_flows)
|
||||
elif key == "v":
|
||||
if self.viewing == self.REQ:
|
||||
conn = self.flow.request
|
||||
@ -519,7 +499,7 @@ class ConsoleFlow(flow.Flow):
|
||||
class ConsoleState(flow.State):
|
||||
def __init__(self):
|
||||
flow.State.__init__(self)
|
||||
self.focus = False
|
||||
self.focus = None
|
||||
self.beep = None
|
||||
|
||||
def add_browserconnect(self, f):
|
||||
@ -545,13 +525,6 @@ class ConsoleState(flow.State):
|
||||
self.set_focus(self.focus)
|
||||
return ret
|
||||
|
||||
@property
|
||||
def view(self):
|
||||
if self.limit:
|
||||
return [i for i in self.flow_list if i.match(self.limit)]
|
||||
else:
|
||||
return self.flow_list[:]
|
||||
|
||||
def get_focus(self):
|
||||
if not self.view or self.focus is None:
|
||||
return None, None
|
||||
@ -580,8 +553,6 @@ class ConsoleState(flow.State):
|
||||
return self.get_from_pos(pos-1)
|
||||
|
||||
def delete_flow(self, f):
|
||||
if not f.intercepting:
|
||||
self.view[self.focus].focus = False
|
||||
ret = flow.State.delete_flow(self, f)
|
||||
self.set_focus(self.focus)
|
||||
return ret
|
||||
@ -694,6 +665,33 @@ class ConsoleMaster(controller.Master):
|
||||
self.nested = True
|
||||
self.make_view()
|
||||
|
||||
def save_flows(self, path):
|
||||
if not path:
|
||||
return
|
||||
data = self.state.dump_flows()
|
||||
path = os.path.expanduser(path)
|
||||
try:
|
||||
f = file(path, "wb")
|
||||
f.write(data)
|
||||
f.close()
|
||||
except IOError, v:
|
||||
self.statusbar.message(str(v))
|
||||
|
||||
def load_flows(self, path):
|
||||
if not path:
|
||||
return
|
||||
path = os.path.expanduser(path)
|
||||
try:
|
||||
f = file(path, "r")
|
||||
data = f.read()
|
||||
f.close()
|
||||
except IOError, v:
|
||||
self.statusbar.message(str(v))
|
||||
return
|
||||
self.state.load_flows(data, ConsoleFlow)
|
||||
self.conn_list_view.set_focus(0)
|
||||
self.sync_list_view()
|
||||
|
||||
def helptext(self):
|
||||
text = []
|
||||
text.extend([("head", "Global keys:\n")])
|
||||
@ -701,13 +699,15 @@ class ConsoleMaster(controller.Master):
|
||||
("A", "accept all intercepted connections"),
|
||||
("a", "accept this intercepted connection"),
|
||||
("B", "set beep filter pattern"),
|
||||
("c", "set sticky cookie expression"),
|
||||
("i", "set interception pattern"),
|
||||
("j, k", "up, down"),
|
||||
("l", "set limit filter pattern"),
|
||||
("L", "load saved flows"),
|
||||
("q", "quit / return to connection list"),
|
||||
("r", "replay request"),
|
||||
("s", "set sticky cookie expression"),
|
||||
("R", "revert changes to request"),
|
||||
("S", "save flows matching current limit"),
|
||||
("page up/down", "page up/down"),
|
||||
("space", "page down"),
|
||||
("enter", "view connection"),
|
||||
@ -923,7 +923,13 @@ class ConsoleMaster(controller.Master):
|
||||
self.view_connlist()
|
||||
else:
|
||||
raise Stop
|
||||
elif k == "s":
|
||||
elif k == "S":
|
||||
self.prompt("Save flows: ", self.save_flows)
|
||||
k = None
|
||||
elif k == "L":
|
||||
self.prompt("Load flows: ", self.load_flows)
|
||||
k = None
|
||||
elif k == "c":
|
||||
self.prompt("Sticky cookie: ", self.set_stickycookie)
|
||||
k = None
|
||||
if k:
|
||||
|
||||
@ -2,6 +2,7 @@
|
||||
This module provides more sophisticated flow tracking. These match requests
|
||||
with their responses, and provide filtering and interception facilities.
|
||||
"""
|
||||
import json
|
||||
import proxy, threading
|
||||
|
||||
class ReplayConnection:
|
||||
@ -27,7 +28,6 @@ class Flow:
|
||||
def __init__(self, connection):
|
||||
self.connection = connection
|
||||
self.request, self.response, self.error = None, None, None
|
||||
self.waiting = True
|
||||
self.intercepting = False
|
||||
self._backup = None
|
||||
|
||||
@ -40,7 +40,7 @@ class Flow:
|
||||
|
||||
@classmethod
|
||||
def from_state(klass, state):
|
||||
f = Flow(ReplayConnection)
|
||||
f = klass(None)
|
||||
if state["request"]:
|
||||
f.request = proxy.Request.from_state(state["request"])
|
||||
if state["response"]:
|
||||
@ -63,7 +63,6 @@ class Flow:
|
||||
|
||||
def revert(self):
|
||||
if self._backup:
|
||||
self.waiting = False
|
||||
restore = [i.copy() if i else None for i in self._backup]
|
||||
self.connection, self.request, self.response, self.error = restore
|
||||
|
||||
@ -133,7 +132,6 @@ class State:
|
||||
if not f:
|
||||
return False
|
||||
f.response = resp
|
||||
f.waiting = False
|
||||
f.backup()
|
||||
return f
|
||||
|
||||
@ -146,16 +144,31 @@ class State:
|
||||
if not f:
|
||||
return None
|
||||
f.error = err
|
||||
f.waiting = False
|
||||
f.backup()
|
||||
return f
|
||||
|
||||
def dump_flows(self):
|
||||
data = [i.get_state() for i in self.view]
|
||||
return json.dumps(data)
|
||||
|
||||
def load_flows(self, js, klass):
|
||||
data = json.loads(js)
|
||||
data = [klass.from_state(i) for i in data]
|
||||
self.flow_list.extend(data)
|
||||
|
||||
def set_limit(self, limit):
|
||||
"""
|
||||
Limit is a compiled filter expression, or None.
|
||||
"""
|
||||
self.limit = limit
|
||||
|
||||
@property
|
||||
def view(self):
|
||||
if self.limit:
|
||||
return tuple([i for i in self.flow_list if i.match(self.limit)])
|
||||
else:
|
||||
return tuple(self.flow_list[:])
|
||||
|
||||
def get_connection(self, itm):
|
||||
if isinstance(itm, (proxy.BrowserConnection, ReplayConnection)):
|
||||
return itm
|
||||
@ -176,7 +189,8 @@ class State:
|
||||
def delete_flow(self, f):
|
||||
if not f.intercepting:
|
||||
c = self.get_connection(f)
|
||||
del self.flow_map[c]
|
||||
if c in self.flow_map:
|
||||
del self.flow_map[c]
|
||||
self.flow_list.remove(f)
|
||||
return True
|
||||
return False
|
||||
@ -214,7 +228,8 @@ class State:
|
||||
if f.request:
|
||||
f.backup()
|
||||
conn = self.get_connection(f)
|
||||
del self.flow_map[conn]
|
||||
if conn in self.flow_map:
|
||||
del self.flow_map[conn]
|
||||
rp = ReplayConnection()
|
||||
f.connection = rp
|
||||
f.request.connection = rp
|
||||
|
||||
@ -57,7 +57,6 @@ class uState(libpry.AutoTree):
|
||||
resp = tresp(req)
|
||||
assert c.add_response(resp)
|
||||
assert len(c.flow_list) == 1
|
||||
assert f.waiting == False
|
||||
assert c.lookup(resp)
|
||||
|
||||
newresp = tresp()
|
||||
@ -183,6 +182,20 @@ class uState(libpry.AutoTree):
|
||||
c.clear()
|
||||
assert len(c.flow_list) == 0
|
||||
|
||||
def test_dump_flows(self):
|
||||
c = console.ConsoleState()
|
||||
self._add_request(c)
|
||||
self._add_response(c)
|
||||
self._add_request(c)
|
||||
self._add_response(c)
|
||||
self._add_request(c)
|
||||
self._add_response(c)
|
||||
|
||||
dump = c.dump_flows()
|
||||
c.clear()
|
||||
c.load_flows(dump, console.ConsoleFlow)
|
||||
assert isinstance(c.flow_list[0], console.ConsoleFlow)
|
||||
|
||||
|
||||
class uFlow(libpry.AutoTree):
|
||||
def test_match(self):
|
||||
|
||||
Loading…
Reference in New Issue
Block a user