Adapt examples

3 files changed, 17 insertions, 18 deletions

diff --git a/examples/flowbasic b/examples/flowbasic
index 4a87b86a..d8d55caa 100644..100755
--- a/examples/flowbasic
+++ b/examples/flowbasic
@@ -8,7 +8,7 @@
     Note that request and response messages are not automatically replied to,
     so we need to implement handlers to do this.
 """
-from mitmproxy import flow
+from mitmproxy import flow, controller
 from mitmproxy.proxy import ProxyServer, ProxyConfig
 
 
@@ -19,18 +19,15 @@ class MyMaster(flow.FlowMaster):
         except KeyboardInterrupt:
             self.shutdown()
 
+    @controller.handler
     def handle_request(self, f):
         f = flow.FlowMaster.handle_request(self, f)
-        if f:
-            f.reply()
-        return f
+        print(f)
 
+    @controller.handler
     def handle_response(self, f):
         f = flow.FlowMaster.handle_response(self, f)
-        if f:
-            f.reply()
         print(f)
-        return f
 
 
 config = ProxyConfig(
diff --git a/examples/stickycookies b/examples/stickycookies
index 8f11de8d..43e5371d 100644..100755
--- a/examples/stickycookies
+++ b/examples/stickycookies
@@ -21,19 +21,19 @@ class StickyMaster(controller.Master):
         except KeyboardInterrupt:
             self.shutdown()
 
-    def handle_request(self, flow):
+    @controller.handler
+    def request(self, flow):
         hid = (flow.request.host, flow.request.port)
         if "cookie" in flow.request.headers:
             self.stickyhosts[hid] = flow.request.headers.get_all("cookie")
         elif hid in self.stickyhosts:
             flow.request.headers.set_all("cookie", self.stickyhosts[hid])
-        flow.reply()
 
-    def handle_response(self, flow):
+    @controller.handler
+    def response(self, flow):
         hid = (flow.request.host, flow.request.port)
         if "set-cookie" in flow.response.headers:
             self.stickyhosts[hid] = flow.response.headers.get_all("set-cookie")
-        flow.reply()
 
 
 config = proxy.ProxyConfig(port=8080)
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py
index 51d58723..dcf920ef 100644
--- a/mitmproxy/controller.py
+++ b/mitmproxy/controller.py
@@ -36,10 +36,18 @@ class Master(object):
     """
         The master handles mitmproxy's main event loop.
     """
-    def __init__(self):
+    def __init__(self, *servers):
         self.event_queue = queue.Queue()
         self.should_exit = threading.Event()
         self.servers = []
+        for i in servers:
+            self.add_server(i)
+
+    def add_server(self, server):
+        # We give a Channel to the server which can be used to communicate with the master
+        channel = Channel(self.event_queue, self.should_exit)
+        server.set_channel(channel)
+        self.servers.append(server)
 
     def start(self):
         self.should_exit.clear()
@@ -87,12 +95,6 @@ class Master(object):
             server.shutdown()
         self.should_exit.set()
 
-    def add_server(self, server):
-        # We give a Channel to the server which can be used to communicate with the master
-        channel = Channel(self.event_queue, self.should_exit)
-        server.set_channel(channel)
-        self.servers.append(server)
-
 
 class ServerThread(threading.Thread):
     def __init__(self, server):