remove script contexts

1 files changed, 32 insertions, 14 deletions

diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py
index a170d868..222ebc69 100644
--- a/mitmproxy/controller.py
+++ b/mitmproxy/controller.py
@@ -2,11 +2,12 @@ from __future__ import absolute_import, print_function, division
 
 import functools
 import threading
+import contextlib
 
 from six.moves import queue
 
+import mitmproxy
 from netlib import basethread
-
 from . import exceptions
 
 
@@ -34,6 +35,16 @@ Events = frozenset([
 ])
 
 
+class Log(object):
+    def __init__(self, master):
+        self.master = master
+
+    def __call__(self, text, level="info"):
+        self.master.add_event(text, level)
+
+    # We may want to add .log(), .warn() etc. here at a later point in time
+
+
 class Master(object):
     """
         The master handles mitmproxy's main event loop.
@@ -45,6 +56,20 @@ class Master(object):
         for i in servers:
             self.add_server(i)
 
+    @contextlib.contextmanager
+    def handlecontext(self):
+        # Handlecontexts also have to nest - leave cleanup to the outermost
+        if mitmproxy.master:
+            yield
+            return
+        mitmproxy.master = self
+        mitmproxy.log = Log(self)
+        try:
+            yield
+        finally:
+            mitmproxy.master = None
+            mitmproxy.log = None
+
     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)
@@ -77,8 +102,8 @@ class Master(object):
                 if mtype not in Events:
                     raise exceptions.ControlException("Unknown event %s" % repr(mtype))
                 handle_func = getattr(self, mtype)
-                if not hasattr(handle_func, "__dict__"):
-                    raise exceptions.ControlException("Handler %s not a function" % mtype)
+                if not callable(handle_func):
+                    raise exceptions.ControlException("Handler %s not callable" % mtype)
                 if not handle_func.__dict__.get("__handler"):
                     raise exceptions.ControlException(
                         "Handler function %s is not decorated with controller.handler" % (
@@ -151,15 +176,7 @@ class Channel(object):
 
 def handler(f):
     @functools.wraps(f)
-    def wrapper(*args, **kwargs):
-        # We can either be called as a method, or as a wrapped solo function
-        if len(args) == 1:
-            message = args[0]
-        elif len(args) == 2:
-            message = args[1]
-        else:
-            raise exceptions.ControlException("Handler takes one argument: a message")
-
+    def wrapper(master, message):
         if not hasattr(message, "reply"):
             raise exceptions.ControlException("Message %s has no reply attribute" % message)
 
@@ -172,7 +189,8 @@ def handler(f):
             handling = True
             message.reply.handled = True
 
-        ret = f(*args, **kwargs)
+        with master.handlecontext():
+            ret = f(master, message)
 
         if handling and not message.reply.acked and not message.reply.taken:
             message.reply.ack()
@@ -216,7 +234,7 @@ class Reply(object):
     def __del__(self):
         if not self.acked:
             # This will be ignored by the interpreter, but emit a warning
-            raise exceptions.ControlException("Un-acked message")
+            raise exceptions.ControlException("Un-acked message: %s" % self.obj)
 
 
 class DummyReply(object):