Add IRCv3 cap support

Author: linuxdaemon

cloudbot/clients/irc.py

@@ -122,7 +122,8 @@ 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
+        # send the cap ls, password, nick, and user
+        self.send("CAP LS 302")
         self.set_pass(self.config["connection"].get("password"))
         self.set_nick(self.nick)
         self.cmd("USER", self.config.get('user', 'cloudbot'), "3", "*",

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,37 @@ 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

cloudbot/plugin.py

@@ -5,6 +5,7 @@
 import logging
 import os
 import re
+from collections import defaultdict
 
 import sqlalchemy
 
@@ -18,7 +19,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])
     """
     # set the loaded flag
     module._cloudbot_loaded = True
@@ -30,8 +31,11 @@ def find_hooks(parent, module):
     periodic = []
     on_start = []
     on_stop = []
+    on_cap_ack = []
+    on_cap_available = []
     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}
     for name, func in module.__dict__.items():
         if hasattr(func, "_cloudbot_hook"):
             # if it has cloudbot hook
@@ -43,7 +47,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
 
 
 def find_tables(code):
@@ -98,6 +102,7 @@ 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._hook_waiting_queues = {}
 
     @asyncio.coroutine
@@ -179,6 +184,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)
@@ -259,6 +274,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:
@@ -521,7 +552,8 @@ 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)
+        self.commands, self.regexes, self.raw_hooks, self.sieves, self.events, self.periodic, self.run_on_start, self.run_on_stop, self.on_cap_ack, self.on_cap_available = find_hooks(
+            self, code)
         # 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 +808,28 @@ 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)
+
+
 _hook_name_to_plugin = {
     "command": CommandHook,
     "regex": RegexHook,
@@ -785,4 +839,6 @@ def __str__(self):
     "periodic": PeriodicHook,
     "on_start": OnStartHook,
     "on_stop": OnStopHook,
+    "on_cap_available": OnCapAvaliableHook,
+    "on_cap_ack": OnCapAckHook,
 }

plugins/cap.py

@@ -0,0 +1,75 @@
+import asyncio
+import logging
+from functools import partial
+
+from cloudbot import hook
+from cloudbot.event import CapEvent
+
+logger = logging.getLogger("cloudbot")
+
+
+@asyncio.coroutine
+def handle_available_caps(conn, caplist, event, irc_paramlist, bot):
+    available_caps = conn.memory.setdefault("available_caps", set())
+    caps = [tuple(cap.split('=', 1)) for cap in caplist]
+    available_caps.update(caps)
+    cap_queue = conn.memory.setdefault("cap_queue", {})
+    for cap, *param in caps:
+        cap_event = partial(CapEvent, base_event=event, cap=cap, cap_param=param[0] if param else None)
+        tasks = [
+            bot.plugin_manager.launch(_hook, cap_event(hook=_hook))
+            for _hook in bot.plugin_manager.cap_hooks["on_available"][cap.casefold()]
+        ]
+        results = yield from asyncio.gather(*tasks)
+        if any(results):
+            cap_queue[cap.casefold()] = conn.loop.create_future()
+            conn.cmd("CAP", "REQ", cap)
+
+    if irc_paramlist[2] != '+':
+        yield from asyncio.gather(*cap_queue.values())
+        cap_queue.clear()
+        conn.send("CAP END")
+
+
+@asyncio.coroutine
+@hook.irc_raw("CAP")
+def on_cap(irc_paramlist, conn, bot, event):
+    caplist = []
+    if len(irc_paramlist) > 2:
+        capstr = irc_paramlist[-1].strip()
+        if capstr[0] == ':':
+            capstr = capstr[1:]
+
+        caplist = capstr.split()
+    subcmd = irc_paramlist[1].upper()
+    if subcmd == "LS":
+        yield from handle_available_caps(conn, caplist, event, irc_paramlist, bot)
+
+    elif subcmd in ('ACK', 'NAK'):
+        enabled = subcmd == 'ACK'
+        server_caps = conn.memory.setdefault('server_caps', {})
+        cap_queue = conn.memory.get("cap_queue", {})
+        caps = [cap.casefold() for cap in caplist]
+        for cap in caps:
+            server_caps[cap] = enabled
+            if enabled:
+                cap_event = partial(CapEvent, base_event=event, cap=cap)
+                tasks = [
+                    bot.plugin_manager.launch(_hook, cap_event(hook=_hook))
+                    for _hook in bot.plugin_manager.cap_hooks["on_ack"][cap]
+                ]
+                yield from asyncio.gather(*tasks)
+
+            if cap in cap_queue:
+                cap_queue[cap].set_result(enabled)
+
+    elif subcmd == 'LIST':
+        logger.info("Enabled Capabilities: %s", irc_paramlist[-1])
+    elif subcmd == 'NEW':
+        logger.info("New capabilities advertised: %s", irc_paramlist[-1])
+        yield from handle_available_caps(conn, caplist, event, irc_paramlist, bot)
+    elif subcmd == 'DEL':
+        logger.info("Capabilities removed by server: %s", irc_paramlist[-1])
+        server_caps = conn.memory.setdefault('server_caps', {})
+        for cap in caplist:
+            server_caps[cap] = False