Merge pull request #1626 from cortesi/state

Convert flow.state to an addon

12 files changed, 37 insertions, 34 deletions

diff --git a/mitmproxy/addons.py b/mitmproxy/addons.py
index 5b078129..9e5677da 100644
--- a/mitmproxy/addons.py
+++ b/mitmproxy/addons.py
@@ -25,7 +25,7 @@ class Addons:
     def _options_update(self, options, updated):
         for i in self.chain:
             with self.master.handlecontext():
-                i.configure(options, updated)
+                self.invoke_with_context(i, "configure", options, updated)
 
     def startup(self, s):
         """
diff --git a/mitmproxy/console/flowlist.py b/mitmproxy/console/flowlist.py
index 653dfa02..ba555647 100644
--- a/mitmproxy/console/flowlist.py
+++ b/mitmproxy/console/flowlist.py
@@ -158,7 +158,7 @@ class ConnectionItem(urwid.WidgetWrap):
         (maxcol,) = xxx_todo_changeme
         key = common.shortcuts(key)
         if key == "a":
-            self.flow.accept_intercept(self.master)
+            self.flow.resume(self.master)
             signals.flowlist_change.send(self)
         elif key == "d":
             if self.flow.killable:
diff --git a/mitmproxy/console/flowview.py b/mitmproxy/console/flowview.py
index 0422e72b..b8f91bdb 100644
--- a/mitmproxy/console/flowview.py
+++ b/mitmproxy/console/flowview.py
@@ -519,7 +519,7 @@ class FlowView(tabs.Tabs):
             # Pass scroll events to the wrapped widget
             self._w.keypress(size, key)
         elif key == "a":
-            self.flow.accept_intercept(self.master)
+            self.flow.resume(self.master)
             signals.flow_change.send(self, flow = self.flow)
         elif key == "A":
             self.master.accept_all()
diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py
index 1a413990..64761e40 100644
--- a/mitmproxy/console/master.py
+++ b/mitmproxy/console/master.py
@@ -222,7 +222,9 @@ class ConsoleMaster(flow.FlowMaster):
     palette = []
 
     def __init__(self, server, options):
-        flow.FlowMaster.__init__(self, options, server, ConsoleState())
+        flow.FlowMaster.__init__(self, options, server)
+        self.state = ConsoleState()
+        self.addons.add(self.state)
         self.stream_path = None
         # This line is just for type hinting
         self.options = self.options  # type: Options
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py
index 148fda77..a225634a 100644
--- a/mitmproxy/controller.py
+++ b/mitmproxy/controller.py
@@ -27,6 +27,9 @@ Events = frozenset([
     "responseheaders",
     "error",
 
+    "intercept",
+    "resume",
+
     "websocket_handshake",
 
     "next_layer",
diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py
index d3f894fa..d3a66876 100644
--- a/mitmproxy/dump.py
+++ b/mitmproxy/dump.py
@@ -33,7 +33,7 @@ class Options(options.Options):
 class DumpMaster(flow.FlowMaster):
 
     def __init__(self, server, options):
-        flow.FlowMaster.__init__(self, options, server, flow.DummyState())
+        flow.FlowMaster.__init__(self, options, server)
         self.has_errored = False
         self.addons.add(termlog.TermLog())
         self.addons.add(*builtins.default_addons())
diff --git a/mitmproxy/flow/__init__.py b/mitmproxy/flow/__init__.py
index 9d243f49..07e1de6c 100644
--- a/mitmproxy/flow/__init__.py
+++ b/mitmproxy/flow/__init__.py
@@ -4,10 +4,10 @@ from mitmproxy.flow.master import FlowMaster
 from mitmproxy.flow.modules import (
     AppRegistry
 )
-from mitmproxy.flow.state import State, DummyState, FlowView
+from mitmproxy.flow.state import State, FlowView
 
 __all__ = [
     "export", "modules",
     "FlowWriter", "FilteredFlowWriter", "FlowReader", "read_flows_from_paths",
-    "FlowMaster", "AppRegistry", "DummyState", "State", "FlowView",
+    "FlowMaster", "AppRegistry", "State", "FlowView",
 ]
diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py
index 3d57c7bd..33502dc5 100644
--- a/mitmproxy/flow/master.py
+++ b/mitmproxy/flow/master.py
@@ -47,11 +47,10 @@ class FlowMaster(controller.Master):
         if len(self.servers) > 0:
             return self.servers[0]
 
-    def __init__(self, options, server, state):
+    def __init__(self, options, server):
         super().__init__(options)
         if server:
             self.add_server(server)
-        self.state = state
         self.apps = modules.AppRegistry()
 
     def start_app(self, host, port):
@@ -196,7 +195,7 @@ class FlowMaster(controller.Master):
 
     @controller.handler
     def error(self, f):
-        self.state.update_flow(f)
+        pass
 
     @controller.handler
     def requestheaders(self, f):
@@ -216,8 +215,6 @@ class FlowMaster(controller.Master):
                     self.add_log("Error in wsgi app. %s" % err, "error")
                 f.reply.kill()
                 return
-        if f not in self.state.flows:  # don't add again on replay
-            self.state.add_flow(f)
 
     @controller.handler
     def responseheaders(self, f):
@@ -225,22 +222,14 @@ class FlowMaster(controller.Master):
 
     @controller.handler
     def response(self, f):
-        self.state.update_flow(f)
+        pass
 
     @controller.handler
     def websocket_handshake(self, f):
         pass
 
-    def handle_intercept(self, f):
-        self.state.update_flow(f)
-
-    def handle_accept_intercept(self, f):
-        self.state.update_flow(f)
-
     @controller.handler
     def tcp_start(self, flow):
-        # TODO: This would break mitmproxy currently.
-        # self.state.add_flow(flow)
         pass
 
     @controller.handler
diff --git a/mitmproxy/flow/state.py b/mitmproxy/flow/state.py
index 8c8e75c7..bb7460b6 100644
--- a/mitmproxy/flow/state.py
+++ b/mitmproxy/flow/state.py
@@ -167,7 +167,7 @@ class FlowStore(FlowList):
     # TODO: Should accept_all operate on views or on all flows?
     def accept_all(self, master):
         for f in self._list:
-            f.accept_intercept(master)
+            f.resume(master)
 
     def kill_all(self, master):
         for f in self._list:
@@ -270,12 +270,19 @@ class State:
         self.add_flow(f2)
         return f2
 
+    # Event handlers
+    def intercept(self, f):
+        self.update_flow(f)
 
-class DummyState:
-    flows = ()
+    def resume(self, f):
+        self.update_flow(f)
 
-    def add_flow(self, *args, **kwargs):
-        pass
+    def error(self, f):
+        self.update_flow(f)
 
-    def update_flow(self, *args, **kwargs):
-        pass
+    def request(self, f):
+        if f not in self.flows:  # don't add again on replay
+            self.add_flow(f)
+
+    def response(self, f):
+        self.update_flow(f)
diff --git a/mitmproxy/models/flow.py b/mitmproxy/models/flow.py
index 6d3fa0d5..2596165b 100644
--- a/mitmproxy/models/flow.py
+++ b/mitmproxy/models/flow.py
@@ -170,16 +170,16 @@ class Flow(stateobject.StateObject):
 
     def intercept(self, master):
         """
-            Intercept this Flow. Processing will stop until accept_intercept is
+            Intercept this Flow. Processing will stop until resume is
             called.
         """
         if self.intercepted:
             return
         self.intercepted = True
         self.reply.take()
-        master.handle_intercept(self)
+        master.addons("intercept", self)
 
-    def accept_intercept(self, master):
+    def resume(self, master):
         """
             Continue with the flow - called after an intercept().
         """
@@ -188,4 +188,4 @@ class Flow(stateobject.StateObject):
         self.intercepted = False
         self.reply.ack()
         self.reply.commit()
-        master.handle_accept_intercept(self)
+        master.addons("intercept", self)
diff --git a/mitmproxy/web/app.py b/mitmproxy/web/app.py
index 46bdd9e3..a81d04be 100644
--- a/mitmproxy/web/app.py
+++ b/mitmproxy/web/app.py
@@ -224,7 +224,7 @@ class AcceptFlows(RequestHandler):
 class AcceptFlow(RequestHandler):
 
     def post(self, flow_id):
-        self.flow.accept_intercept(self.master)
+        self.flow.resume(self.master)
 
 
 class FlowHandler(RequestHandler):
diff --git a/mitmproxy/web/master.py b/mitmproxy/web/master.py
index 23c95e57..3f61fd94 100644
--- a/mitmproxy/web/master.py
+++ b/mitmproxy/web/master.py
@@ -133,7 +133,9 @@ class Options(options.Options):
 class WebMaster(flow.FlowMaster):
 
     def __init__(self, server, options):
-        super().__init__(options, server, WebState())
+        super().__init__(options, server)
+        self.state = WebState()
+        self.addons.add(self.state)
         self.addons.add(*builtins.default_addons())
         self.app = app.Application(
             self, self.options.wdebug, self.options.wauthenticator