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")
if not hasattr(message, "reply"):
raise exceptions.ControlException(
"Message %s has no reply attribute" % message)
# The following ensures that inheritance with wrapped handlers in the
# base class works. If we're the first handler, then responsibility for
# acking is ours. If not, it's someone else's and we ignore it.
handling = False
# We're the first handler - ack responsibility is ours
if not message.reply.handled:
handling = True
message.reply.handled = True
ret = f(*args, **kwargs)
if handling and not message.reply.acked and not message.reply.taken:
message.reply.ack()
return ret
def tick(self, timeout):
if self.first_tick:
self.first_tick = False
self.addons.invoke_all_with_context("running")
with self.handlecontext():
self.addons("tick")
changed = False
try:
mtype, obj = self.event_queue.get(timeout=timeout)
if mtype not in eventsequence.Events:
raise exceptions.ControlException("Unknown event %s" %
repr(mtype))
handle_func = getattr(self, 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"
% (handle_func))
handle_func(obj)
self.event_queue.task_done()
changed = True
except queue.Empty:
pass
return changed
def send(self, msg, force=False):
if self.state not in ("handled", "taken"):
raise exceptions.ControlException(
"Reply is {}, did not expect a call to .send().".format(self.state)
)
if self.has_message and not force:
raise exceptions.ControlException("There is already a reply message.")
self.value = msg
def send(self, msg, force=False):
if self.state not in {"start", "taken"}:
raise exceptions.ControlException(
"Reply is {}, but expected it to be start or taken.".format(self.state)
)
if self.has_message and not force:
raise exceptions.ControlException("There is already a reply message.")
self.value = msg
def commit(self):
"""
Ultimately, messages are commited. This is done either automatically by the handler
if the message is not taken or manually by the entity which called .take().
"""
if self.state != "taken":
raise exceptions.ControlException("Reply is {}, but expected it to be taken.".format(self.state))
if not self.has_message:
raise exceptions.ControlException("There is no reply message.")
self._state = "committed"
self.q.put(self.value)
async def handle_lifecycle(self, name, message):
"""
Handle a lifecycle event.
"""
if not hasattr(message, "reply"): # pragma: no cover
raise exceptions.ControlException(
"Message %s has no reply attribute" % message
)
# We can use DummyReply objects multiple times. We only clear them up on
# the next handler so that we can access value and state in the
# meantime.
if isinstance(message.reply, controller.DummyReply):
message.reply.reset()
self.trigger(name, message)
if message.reply.state == "start":
message.reply.take()
if not message.reply.has_message:
message.reply.ack()
message.reply.commit()
if isinstance(message.reply, controller.DummyReply):
message.reply.mark_reset()
if isinstance(message, flow.Flow):
self.trigger("update", [message])
def take(self):
"""
Scripts or other parties make "take" a reply out of a normal flow.
For example, intercepted flows are taken out so that the connection thread does not proceed.
"""
if self.state != "handled":
raise exceptions.ControlException("Reply is {}, but expected it to be handled.".format(self.state))
self._state = "taken"
def handle(self):
"""
Reply are handled by controller.handlers, which may be nested. The first handler takes
responsibility and handles the reply.
"""
if self.state != "unhandled":
raise exceptions.ControlException("Reply is {}, but expected it to be unhandled.".format(self.state))
self._state = "handled"
def __call__(self, msg=NO_REPLY):
if self.acked:
raise exceptions.ControlException("Message already acked.")
self.acked = True
if msg is NO_REPLY:
self.q.put(self.obj)
else:
self.q.put(msg)
def kill(self):
"""
Kill this flow. The current request/response will not be forwarded to its destination.
"""
if not self.killable:
raise exceptions.ControlException("Flow is not killable.")
self.error = Error(Error.KILLED_MESSAGE)
self.intercepted = False
self.live = False
def kill(self):
"""
Kill this request.
"""
if not self.killable:
raise exceptions.ControlException("Flow is not killable.")
self.error = Error(Error.KILLED_MESSAGE)
self.intercepted = False
self.live = False
async def main(self):
while True:
try:
mtype, obj = await self.event_queue.get()
except RuntimeError:
return
if mtype not in eventsequence.Events: # pragma: no cover
raise exceptions.ControlException("Unknown event %s" % repr(mtype))
self.addons.handle_lifecycle(mtype, obj)
self.event_queue.task_done()
def kill(self):
"""
Kill this flow. The current request/response will not be forwarded to its destination.
"""
if not self.killable:
raise exceptions.ControlException("Flow is not killable.")
# TODO: The way we currently signal killing is not ideal. One major problem is that we cannot kill
# flows in transit (https://github.com/mitmproxy/mitmproxy/issues/4711), even though they are advertised
# as killable. An alternative approach would be to introduce a `KillInjected` event similar to
# `MessageInjected`, which should fix this issue.
self.error = Error(Error.KILLED_MESSAGE)
self.intercepted = False
self.live = False
def commit(self):
"""
Ultimately, messages are committed. This is done either automatically by
the handler if the message is not taken or manually by the entity which
called .take().
"""
if self.state != "taken":
raise exceptions.ControlException(
f"Reply is {self.state}, but expected it to be taken.")
self._state = "committed"
try:
self._loop.call_soon_threadsafe(lambda: self.done.set())
except RuntimeError: # pragma: no cover
pass # event loop may already be closed.
def tick(self, timeout):
changed = False
try:
# This endless loop runs until the 'Queue.Empty'
# exception is thrown.
while True:
mtype, obj = self.event_queue.get(timeout=timeout)
if mtype not in Events:
raise exceptions.ControlException("Unknown event %s" %
repr(mtype))
handle_func = getattr(self, mtype)
if not hasattr(handle_func, "func_dict"):
raise exceptions.ControlException(
"Handler %s not a function" % mtype)
if not handle_func.func_dict.get("__handler"):
raise exceptions.ControlException(
"Handler function %s is not decorated with controller.handler"
% (handle_func))
handle_func(obj)
self.event_queue.task_done()
changed = True
except queue.Empty:
pass
return changed
def tick(self, timeout):
if self.first_tick:
self.first_tick = False
self.addons.trigger("running")
self.addons.trigger("tick")
changed = False
try:
mtype, obj = self.event_queue.get(timeout=timeout)
if mtype not in eventsequence.Events:
raise exceptions.ControlException("Unknown event %s" %
repr(mtype))
self.addons.handle_lifecycle(mtype, obj)
self.event_queue.task_done()
changed = True
except queue.Empty:
pass
return changed
def wrapper(master, message):
if not hasattr(message, "reply"):
raise exceptions.ControlException(
"Message %s has no reply attribute" % message)
# DummyReplys may be reused multiple times.
# We only clear them up on the next handler so that we can access value and
# state in the meantime.
if isinstance(message.reply, DummyReply):
message.reply.reset()
# The following ensures that inheritance with wrapped handlers in the
# base class works. If we're the first handler, then responsibility for
# acking is ours. If not, it's someone else's and we ignore it.
handling = False
# We're the first handler - ack responsibility is ours
if message.reply.state == "unhandled":
handling = True
message.reply.handle()
with master.handlecontext():
ret = f(master, message)
if handling:
master.addons(f.__name__, message)
# Reset the handled flag - it's common for us to feed the same object
# through handlers repeatedly, so we don't want this to persist across
# calls.
if handling and message.reply.state == "handled":
message.reply.take()
if not message.reply.has_message:
message.reply.ack()
message.reply.commit()
# DummyReplys may be reused multiple times.
if isinstance(message.reply, DummyReply):
message.reply.mark_reset()
return ret
def mark_reset(self):
if self.state != "committed":
raise exceptions.ControlException(f"Uncommitted reply: {self.obj}")
self._should_reset = True
def __del__(self):
if self.state != "committed":
# This will be ignored by the interpreter, but emit a warning
raise exceptions.ControlException(f"Uncommitted reply: {self.obj}")
def __del__(self):
if not self.acked:
# This will be ignored by the interpreter, but emit a warning
raise exceptions.ControlException("Un-acked message")
def send(self, msg):
if self.acked:
raise exceptions.ControlException("Message already acked.")
self.acked = True
self.q.put(msg)
def ack(self, force=False):
if self.state not in {"start", "taken"}:
raise exceptions.ControlException(
"Reply is {}, but expected it to be start or taken.".format(self.state)
)
self.send(self.obj, force)
def send(self, msg, force=False):
if self.has_message and not force:
raise exceptions.ControlException("There is already a reply message.")
self.value = msg
