diff options
Diffstat (limited to 'libmproxy/flow.py')
| -rw-r--r-- | libmproxy/flow.py | 29 |
1 files changed, 22 insertions, 7 deletions
diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 41064829..31b7bc51 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -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 |