Merge pull request CloudBotIRC#146 from linuxdaemon/gonzobot+cap-support

IRCv3 CAP support + on_connect hooks

Author: edwardslabs

cloudbot/clients/irc.py

@@ -1,8 +1,8 @@
-from _ssl import PROTOCOL_SSLv23
 import asyncio
+import logging
 import re
 import ssl
-import logging
+from _ssl import PROTOCOL_SSLv23
 from ssl import SSLContext
 
 from cloudbot.client import Client
@@ -122,11 +122,12 @@ def connect(self):
         self._transport, self._protocol = yield from self.loop.create_connection(
             lambda: _IrcProtocol(self), host=self.server, port=self.port, ssl=self.ssl_context, **optional_params)
 
-        # send the password, nick, and user
-        self.set_pass(self.config["connection"].get("password"))
-        self.set_nick(self.nick)
-        self.cmd("USER", self.config.get('user', 'cloudbot'), "3", "*",
-                 self.config.get('realname', 'CloudBot - https://git.io/CloudBot'))
+        tasks = [
+            self.bot.plugin_manager.launch(hook, Event(bot=self.bot, conn=self, hook=hook))
+            for hook in self.bot.plugin_manager.connect_hooks
+        ]
+        # TODO stop connecting if a connect hook fails?
+        yield from asyncio.gather(*tasks)
 
     def quit(self, reason=None):
         if self._quit:

cloudbot/event.py

@@ -1,7 +1,7 @@
 import asyncio
+import concurrent.futures
 import enum
 import logging
-import concurrent.futures
 
 logger = logging.getLogger("cloudbot")
 
@@ -383,3 +383,10 @@ def __init__(self, *, bot=None, hook, match, conn=None, base_event=None, event_t
                          content_raw=content_raw, target=target, channel=channel, nick=nick, user=user, host=host, mask=mask,
                          irc_raw=irc_raw, irc_prefix=irc_prefix, irc_command=irc_command, irc_paramlist=irc_paramlist)
         self.match = match
+
+
+class CapEvent(Event):
+    def __init__(self, *args, cap, cap_param=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.cap = cap
+        self.cap_param = cap_param

cloudbot/hook.py

@@ -1,6 +1,6 @@
+import collections
 import inspect
 import re
-import collections
 
 from cloudbot.event import EventType
 
@@ -177,6 +177,16 @@ def add_hook(self, trigger_param, kwargs):
             self.types.update(trigger_param)
 
 
+class _CapHook(_Hook):
+    def __init__(self, func, _type):
+        super().__init__(func, "on_cap_{}".format(_type))
+        self.caps = set()
+
+    def add_hook(self, caps, kwargs):
+        self._add_hook(kwargs)
+        self.caps.update(caps)
+
+
 def _add_hook(func, hook):
     if not hasattr(func, "_cloudbot_hook"):
         func._cloudbot_hook = {}
@@ -357,3 +367,55 @@ def _on_stop_hook(func):
         return lambda func: _on_stop_hook(func)
 
 on_unload = on_stop
+
+
+def on_cap_available(*caps, **kwargs):
+    """External on_cap_available decorator. Must be used as a function that returns a decorator
+
+    This hook will fire for each capability in a `CAP LS` response from the server
+    """
+
+    def _on_cap_available_hook(func):
+        hook = _get_hook(func, "on_cap_available")
+        if hook is None:
+            hook = _CapHook(func, "available")
+            _add_hook(func, hook)
+        hook.add_hook(caps, kwargs)
+        return func
+
+    return _on_cap_available_hook
+
+
+def on_cap_ack(*caps, **kwargs):
+    """External on_cap_ack decorator. Must be used as a function that returns a decorator
+
+    This hook will fire for each capability that is acknowledged from the server with `CAP ACK`
+    """
+
+    def _on_cap_ack_hook(func):
+        hook = _get_hook(func, "on_cap_ack")
+        if hook is None:
+            hook = _CapHook(func, "ack")
+            _add_hook(func, hook)
+        hook.add_hook(caps, kwargs)
+        return func
+
+    return _on_cap_ack_hook
+
+
+def on_connect(param=None, **kwargs):
+    def _on_connect_hook(func):
+        hook = _get_hook(func, "on_connect")
+        if hook is None:
+            hook = _Hook(func, "on_connect")
+            _add_hook(func, hook)
+        hook._add_hook(kwargs)
+        return func
+
+    if callable(param):
+        return _on_connect_hook(param)
+    else:
+        return lambda func: _on_connect_hook(func)
+
+
+connect = on_connect

cloudbot/plugin.py

@@ -5,6 +5,8 @@
 import logging
 import os
 import re
+from collections import defaultdict
+from operator import attrgetter
 
 import sqlalchemy
 
@@ -18,7 +20,7 @@ def find_hooks(parent, module):
     """
     :type parent: Plugin
     :type module: object
-    :rtype: (list[CommandHook], list[RegexHook], list[RawHook], list[SieveHook], List[EventHook], List[PeriodicHook], list[OnStartHook], List[OnStopHook])
+    :rtype: (list[CommandHook], list[RegexHook], list[RawHook], list[SieveHook], List[EventHook], List[PeriodicHook], list[OnStartHook], List[OnStopHook], list[OnCapAckHook], list[OnCapAvailableHook], list[OnConnectHook])
     """
     # set the loaded flag
     module._cloudbot_loaded = True
@@ -30,8 +32,12 @@ def find_hooks(parent, module):
     periodic = []
     on_start = []
     on_stop = []
+    on_cap_ack = []
+    on_cap_available = []
+    on_connect = []
     type_lists = {"command": command, "regex": regex, "irc_raw": raw, "sieve": sieve, "event": event,
-                  "periodic": periodic, "on_start": on_start, "on_stop": on_stop}
+                  "periodic": periodic, "on_start": on_start, "on_stop": on_stop, "on_cap_ack": on_cap_ack,
+                  "on_cap_available": on_cap_available, "on_connect": on_connect}
     for name, func in module.__dict__.items():
         if hasattr(func, "_cloudbot_hook"):
             # if it has cloudbot hook
@@ -43,7 +49,7 @@ def find_hooks(parent, module):
             # delete the hook to free memory
             del func._cloudbot_hook
 
-    return command, regex, raw, sieve, event, periodic, on_start, on_stop
+    return command, regex, raw, sieve, event, periodic, on_start, on_stop, on_cap_ack, on_cap_available, on_connect
 
 
 def find_tables(code):
@@ -98,6 +104,8 @@ def __init__(self, bot):
         self.event_type_hooks = {}
         self.regex_hooks = []
         self.sieves = []
+        self.cap_hooks = {"on_available": defaultdict(list), "on_ack": defaultdict(list)}
+        self.connect_hooks = []
         self._hook_waiting_queues = {}
 
     @asyncio.coroutine
@@ -179,6 +187,16 @@ def load_plugin(self, path):
 
         self.plugins[plugin.file_name] = plugin
 
+        for on_cap_available_hook in plugin.on_cap_available:
+            for cap in on_cap_available_hook.caps:
+                self.cap_hooks["on_available"][cap.casefold()].append(on_cap_available_hook)
+            self._log_hook(on_cap_available_hook)
+
+        for on_cap_ack_hook in plugin.on_cap_ack:
+            for cap in on_cap_ack_hook.caps:
+                self.cap_hooks["on_ack"][cap.casefold()].append(on_cap_ack_hook)
+            self._log_hook(on_cap_ack_hook)
+
         for periodic_hook in plugin.periodic:
             task = asyncio.async(self._start_periodic(periodic_hook))
             plugin.tasks.append(task)
@@ -227,8 +245,14 @@ def load_plugin(self, path):
             self.sieves.append(sieve_hook)
             self._log_hook(sieve_hook)
 
+        # register connect hooks
+        for connect_hook in plugin.connect_hooks:
+            self.connect_hooks.append(connect_hook)
+            self._log_hook(connect_hook)
+
         # sort sieve hooks by priority
         self.sieves.sort(key=lambda x: x.priority)
+        self.connect_hooks.sort(key=attrgetter("priority"))
 
         # we don't need this anymore
         del plugin.run_on_start
@@ -259,6 +283,22 @@ def unload_plugin(self, path):
         for task in plugin.tasks:
             task.cancel()
 
+        for on_cap_available_hook in plugin.on_cap_available:
+            available_hooks = self.cap_hooks["on_available"]
+            for cap in on_cap_available_hook.caps:
+                cap_cf = cap.casefold()
+                available_hooks[cap_cf].remove(on_cap_available_hook)
+                if not available_hooks[cap_cf]:
+                    del available_hooks[cap_cf]
+
+        for on_cap_ack in plugin.on_cap_ack:
+            ack_hooks = self.cap_hooks["on_ack"]
+            for cap in on_cap_ack.caps:
+                cap_cf = cap.casefold()
+                ack_hooks[cap_cf].remove(on_cap_ack)
+                if not ack_hooks[cap_cf]:
+                    del ack_hooks[cap_cf]
+
         # unregister commands
         for command_hook in plugin.commands:
             for alias in command_hook.aliases:
@@ -294,6 +334,10 @@ def unload_plugin(self, path):
         for sieve_hook in plugin.sieves:
             self.sieves.remove(sieve_hook)
 
+        # unregister connect hooks
+        for connect_hook in plugin.connect_hooks:
+            self.connect_hooks.remove(connect_hook)
+
         # Run on_stop hooks
         for on_stop_hook in plugin.run_on_stop:
             event = Event(bot=self.bot, hook=on_stop_hook)
@@ -521,7 +565,12 @@ def __init__(self, filepath, filename, title, code):
         self.file_path = filepath
         self.file_name = filename
         self.title = title
-        self.commands, self.regexes, self.raw_hooks, self.sieves, self.events, self.periodic, self.run_on_start, self.run_on_stop = find_hooks(self, code)
+        # TODO clean up hook lists
+        hooks = find_hooks(self, code)
+        self.commands, self.regexes, self.raw_hooks, *hooks = hooks
+        self.sieves, self.events, self.periodic, *hooks = hooks
+        self.run_on_start, self.run_on_stop, self.on_cap_ack, *hooks = hooks
+        self.on_cap_available, self.connect_hooks, *hooks = hooks
         # we need to find tables for each plugin so that they can be unloaded from the global metadata when the
         # plugin is reloaded
         self.tables = find_tables(code)
@@ -776,6 +825,45 @@ def __str__(self):
         return "on_stop {} from {}".format(self.function_name, self.plugin.file_name)
 
 
+class CapHook(Hook):
+    def __init__(self, _type, plugin, base_hook):
+        self.caps = base_hook.caps
+        super().__init__("on_cap_{}".format(_type), plugin, base_hook)
+
+    def __repr__(self):
+        return "{name}[{caps} {base!r}]".format(name=self.type, caps=self.caps, base=super())
+
+    def __str__(self):
+        return "{name} {func} from {file}".format(name=self.type, func=self.function_name, file=self.plugin.file_name)
+
+
+class OnCapAvaliableHook(CapHook):
+    def __init__(self, plugin, base_hook):
+        super().__init__("available", plugin, base_hook)
+
+
+class OnCapAckHook(CapHook):
+    def __init__(self, plugin, base_hook):
+        super().__init__("ack", plugin, base_hook)
+
+
+class OnConnectHook(Hook):
+    def __init__(self, plugin, sieve_hook):
+        """
+        :type plugin: Plugin
+        :type sieve_hook: cloudbot.util.hook._Hook
+        """
+
+        self.priority = sieve_hook.kwargs.pop("priority", 100)
+        super().__init__("on_connect", plugin, sieve_hook)
+
+    def __repr__(self):
+        return "{name}[{base!r}]".format(name=self.type, base=super())
+
+    def __str__(self):
+        return "{name} {func} from {file}".format(name=self.type, func=self.function_name, file=self.plugin.file_name)
+
+
 _hook_name_to_plugin = {
     "command": CommandHook,
     "regex": RegexHook,
@@ -785,4 +873,7 @@ def __str__(self):
     "periodic": PeriodicHook,
     "on_start": OnStartHook,
     "on_stop": OnStopHook,
+    "on_cap_available": OnCapAvaliableHook,
+    "on_cap_ack": OnCapAckHook,
+    "on_connect": OnConnectHook,
 }

config.default.json

@@ -18,6 +18,12 @@
             ],
             "disabled_commands": [],
             "acls": {},
+            "sasl": {
+                "enabled": false,
+                "mechanism": "PLAIN",
+                "user": "",
+                "pass": ""
+            },
             "nickserv": {
                 "enabled": false,
                 "nickserv_password": "",