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