Merge pull request CloudBotIRC#178 from linuxdaemon/gonzobot+outgoing-sieves

Add outgoing raw line hooks

Author: edwardslabs

cloudbot/clients/irc.py

@@ -6,7 +6,7 @@
 from ssl import SSLContext
 
 from cloudbot.client import Client
-from cloudbot.event import Event, EventType
+from cloudbot.event import Event, EventType, IrcOutEvent
 from cloudbot.util import async_util
 
 logger = logging.getLogger("cloudbot")
@@ -21,8 +21,10 @@
 irc_bad_chars = ''.join([chr(x) for x in list(range(0, 1)) + list(range(4, 32)) + list(range(127, 160))])
 irc_clean_re = re.compile('[{}]'.format(re.escape(irc_bad_chars)))
 
+
 def irc_clean(dirty):
-    return irc_clean_re.sub('',dirty)
+    return irc_clean_re.sub('', dirty)
+
 
 irc_command_to_event_type = {
     "PRIVMSG": EventType.message,
@@ -292,9 +294,39 @@ def send(self, line):
         # make sure we are connected before sending
         if not self._connected:
             yield from self._connected_future
-        line = line[:510] + "\r\n"
-        data = line.encode("utf-8", "replace")
-        self._transport.write(data)
+
+        old_line = line
+        filtered = bool(self.bot.plugin_manager.out_sieves)
+
+        for out_sieve in self.bot.plugin_manager.out_sieves:
+            event = IrcOutEvent(
+                bot=self.bot, hook=out_sieve, conn=self.conn, irc_raw=line
+            )
+
+            ok, new_line = yield from self.bot.plugin_manager.internal_launch(out_sieve, event)
+            if not ok:
+                logger.warning("Error occurred in outgoing sieve, falling back to old behavior")
+                logger.debug("Line was: %s", line)
+                filtered = False
+                break
+
+            line = new_line
+            if line is not None and not isinstance(line, bytes):
+                line = str(line)
+
+            if not line:
+                return
+
+        if not filtered:
+            # No outgoing sieves loaded or one of the sieves errored, fall back to old behavior
+            line = old_line[:510] + "\r\n"
+            line = line.encode("utf-8", "replace")
+
+        if not isinstance(line, bytes):
+            # the line must be encoded before we send it, one of the sieves didn't encode it, fall back to the default
+            line = line.encode("utf-8", "replace")
+
+        self._transport.write(line)
 
     def data_received(self, data):
         self._input_buffer += data
@@ -412,12 +444,3 @@ def data_received(self, data):
 
             # handle the message, async
             async_util.wrap_future(self.bot.process(event), loop=self.loop)
-
-# Channel Commands
-# NOTICE #chan :Text
-# PRIVMSG #chan :Text
-# KICK #chan nick :reason
-# JOIN #chan
-# PART #chan :reason
-# MODE #chan +<modes>
-# INVITE nick :#chan

cloudbot/event.py

@@ -7,6 +7,8 @@
 
 import sys
 
+from cloudbot.util.parsers.irc import Message
+
 logger = logging.getLogger("cloudbot")
 
 
@@ -419,3 +421,43 @@ def __init__(self, *args, cap, cap_param=None, **kwargs):
         super().__init__(*args, **kwargs)
         self.cap = cap
         self.cap_param = cap_param
+
+
+class IrcOutEvent(Event):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.parsed_line = None
+
+    @asyncio.coroutine
+    def prepare(self):
+        yield from super().prepare()
+
+        if "parsed_line" in self.hook.required_args:
+            try:
+                self.parsed_line = Message.parse(self.line)
+            except Exception:
+                logger.exception("Unable to parse line requested by hook %s", self.hook)
+                self.parsed_line = None
+
+    def prepare_threaded(self):
+        super().prepare_threaded()
+
+        if "parsed_line" in self.hook.required_args:
+            try:
+                self.parsed_line = Message.parse(self.line)
+            except Exception:
+                logger.exception("Unable to parse line requested by hook %s", self.hook)
+                self.parsed_line = None
+
+    @property
+    def line(self):
+        return str(self.irc_raw)
+
+
+class PostHookEvent(Event):
+    def __init__(self, *args, launched_hook=None, launched_event=None, result=None, error=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.launched_hook = launched_hook
+        self.launched_event = launched_event
+        self.result = result
+        self.error = error

cloudbot/hook.py

@@ -1,7 +1,6 @@
 import collections
 import inspect
 import re
-import collections
 from enum import Enum, unique, IntEnum
 
 from cloudbot.event import EventType
@@ -384,18 +383,21 @@ def on_stop(param=None, **kwargs):
     """External on_stop decorator. Can be used directly as a decorator, or with args to return a decorator
     :type param: function | None
     """
+
     def _on_stop_hook(func):
         hook = _get_hook(func, "on_stop")
         if hook is None:
             hook = _Hook(func, "on_stop")
             _add_hook(func, hook)
         hook._add_hook(kwargs)
         return func
+
     if callable(param):
         return _on_stop_hook(param)
     else:
         return lambda func: _on_stop_hook(func)
 
+
 on_unload = on_stop
 
 
@@ -451,6 +453,42 @@ def _on_connect_hook(func):
 connect = on_connect
 
 
+def irc_out(param=None, **kwargs):
+    def _decorate(func):
+        hook = _get_hook(func, "irc_out")
+        if hook is None:
+            hook = _Hook(func, "irc_out")
+            _add_hook(func, hook)
+
+        hook._add_hook(kwargs)
+        return func
+
+    if callable(param):
+        return _decorate(param)
+    else:
+        return lambda func: _decorate(func)
+
+
+def post_hook(param=None, **kwargs):
+    """
+    This hook will be fired just after a hook finishes executing
+    """
+
+    def _decorate(func):
+        hook = _get_hook(func, "post_hook")
+        if hook is None:
+            hook = _Hook(func, "post_hook")
+            _add_hook(func, hook)
+
+        hook._add_hook(kwargs)
+        return func
+
+    if callable(param):
+        return _decorate(param)
+    else:
+        return lambda func: _decorate(func)
+
+
 def permission(*perms, **kwargs):
     def _perm_hook(func):
         assert len(inspect.getfullargspec(func).args) == 3, \