Merge pull request CloudBotIRC#179 from linuxdaemon/gonzobot+op-tracking

Add channel op tracking

Author: edwardslabs

cloudbot/hook.py

@@ -207,6 +207,16 @@ def add_hook(self, caps, kwargs):
         self.caps.update(caps)
 
 
+class _PermissionHook(_Hook):
+    def __init__(self, func):
+        super().__init__(func, "perm_check")
+        self.perms = set()
+
+    def add_hook(self, perms, kwargs):
+        self._add_hook(kwargs)
+        self.perms.update(perms)
+
+
 def _add_hook(func, hook):
     if not hasattr(func, "_cloudbot_hook"):
         func._cloudbot_hook = {}
@@ -439,3 +449,18 @@ def _on_connect_hook(func):
 
 
 connect = on_connect
+
+
+def permission(*perms, **kwargs):
+    def _perm_hook(func):
+        assert len(inspect.getfullargspec(func).args) == 3, \
+            "Permission hook has incorrect argument count. Needs params: bot, event, hook"
+        hook = _get_hook(func, "perm_check")
+        if hook is None:
+            hook = _PermissionHook(func)
+            _add_hook(func, hook)
+
+        hook.add_hook(perms, kwargs)
+        return func
+
+    return lambda func: _perm_hook(func)

cloudbot/plugin.py

@@ -26,36 +26,23 @@ 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], list[OnCapAckHook], list[OnCapAvailableHook], list[OnConnectHook])
+    :rtype: dict
     """
     # set the loaded flag
     module._cloudbot_loaded = True
-    command = []
-    regex = []
-    raw = []
-    sieve = []
-    event = []
-    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, "on_cap_ack": on_cap_ack,
-                  "on_cap_available": on_cap_available, "on_connect": on_connect}
+    hooks = defaultdict(list)
     for name, func in module.__dict__.items():
         if hasattr(func, "_cloudbot_hook"):
             # if it has cloudbot hook
             func_hooks = func._cloudbot_hook
 
             for hook_type, func_hook in func_hooks.items():
-                type_lists[hook_type].append(_hook_name_to_plugin[hook_type](parent, func_hook))
+                hooks[hook_type].append(_hook_name_to_plugin[hook_type](parent, func_hook))
 
             # delete the hook to free memory
             del func._cloudbot_hook
 
-    return command, regex, raw, sieve, event, periodic, on_start, on_stop, on_cap_ack, on_cap_available, on_connect
+    return hooks
 
 
 def find_tables(code):
@@ -112,6 +99,7 @@ def __init__(self, bot):
         self.sieves = []
         self.cap_hooks = {"on_available": defaultdict(list), "on_ack": defaultdict(list)}
         self.connect_hooks = []
+        self.perm_hooks = defaultdict(list)
         self._hook_waiting_queues = {}
 
     @asyncio.coroutine
@@ -182,7 +170,7 @@ def load_plugin(self, path):
         yield from plugin.create_tables(self.bot)
 
         # run on_start hooks
-        for on_start_hook in plugin.run_on_start:
+        for on_start_hook in plugin.hooks["on_start"]:
             success = yield from self.launch(on_start_hook, Event(bot=self.bot, hook=on_start_hook))
             if not success:
                 logger.warning("Not registering hooks from plugin {}: on_start hook errored".format(plugin.title))
@@ -193,23 +181,23 @@ def load_plugin(self, path):
 
         self.plugins[plugin.file_name] = plugin
 
-        for on_cap_available_hook in plugin.on_cap_available:
+        for on_cap_available_hook in plugin.hooks["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 on_cap_ack_hook in plugin.hooks["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:
+        for periodic_hook in plugin.hooks["periodic"]:
             task = async_util.wrap_future(self._start_periodic(periodic_hook))
             plugin.tasks.append(task)
             self._log_hook(periodic_hook)
 
         # register commands
-        for command_hook in plugin.commands:
+        for command_hook in plugin.hooks["command"]:
             for alias in command_hook.aliases:
                 if alias in self.commands:
                     logger.warning(
@@ -220,7 +208,7 @@ def load_plugin(self, path):
             self._log_hook(command_hook)
 
         # register raw hooks
-        for raw_hook in plugin.raw_hooks:
+        for raw_hook in plugin.hooks["irc_raw"]:
             if raw_hook.is_catch_all():
                 self.catch_all_triggers.append(raw_hook)
             else:
@@ -232,7 +220,7 @@ def load_plugin(self, path):
             self._log_hook(raw_hook)
 
         # register events
-        for event_hook in plugin.events:
+        for event_hook in plugin.hooks["event"]:
             for event_type in event_hook.types:
                 if event_type in self.event_type_hooks:
                     self.event_type_hooks[event_type].append(event_hook)
@@ -241,36 +229,42 @@ def load_plugin(self, path):
             self._log_hook(event_hook)
 
         # register regexps
-        for regex_hook in plugin.regexes:
+        for regex_hook in plugin.hooks["regex"]:
             for regex_match in regex_hook.regexes:
                 self.regex_hooks.append((regex_match, regex_hook))
             self._log_hook(regex_hook)
 
         # register sieves
-        for sieve_hook in plugin.sieves:
+        for sieve_hook in plugin.hooks["sieve"]:
             self.sieves.append(sieve_hook)
             self._log_hook(sieve_hook)
 
         # register connect hooks
-        for connect_hook in plugin.connect_hooks:
+        for connect_hook in plugin.hooks["on_connect"]:
             self.connect_hooks.append(connect_hook)
             self._log_hook(connect_hook)
 
+        for perm_hook in plugin.hooks["perm_check"]:
+            for perm in perm_hook.perms:
+                self.perm_hooks[perm].append(perm_hook)
+
+            self._log_hook(perm_hook)
+
         # sort sieve hooks by priority
         self.sieves.sort(key=lambda x: x.priority)
         self.connect_hooks.sort(key=attrgetter("priority"))
 
         # Sort hooks
         self.regex_hooks.sort(key=lambda x: x[1].priority)
-        dicts_of_lists_of_hooks = (self.event_type_hooks, self.raw_triggers)
+        dicts_of_lists_of_hooks = (self.event_type_hooks, self.raw_triggers, self.perm_hooks)
         lists_of_hooks = [self.catch_all_triggers, self.sieves]
         lists_of_hooks.extend(chain.from_iterable(d.values() for d in dicts_of_lists_of_hooks))
 
         for lst in lists_of_hooks:
             lst.sort(key=lambda x: x.priority)
 
         # we don't need this anymore
-        del plugin.run_on_start
+        del plugin.hooks["on_start"]
 
     @asyncio.coroutine
     def unload_plugin(self, path):
@@ -298,15 +292,15 @@ def unload_plugin(self, path):
         for task in plugin.tasks:
             task.cancel()
 
-        for on_cap_available_hook in plugin.on_cap_available:
+        for on_cap_available_hook in plugin.hooks["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:
+        for on_cap_ack in plugin.hooks["on_cap_ack"]:
             ack_hooks = self.cap_hooks["on_ack"]
             for cap in on_cap_ack.caps:
                 cap_cf = cap.casefold()
@@ -315,14 +309,14 @@ def unload_plugin(self, path):
                     del ack_hooks[cap_cf]
 
         # unregister commands
-        for command_hook in plugin.commands:
+        for command_hook in plugin.hooks["command"]:
             for alias in command_hook.aliases:
                 if alias in self.commands and self.commands[alias] == command_hook:
                     # we need to make sure that there wasn't a conflict, so we don't delete another plugin's command
                     del self.commands[alias]
 
         # unregister raw hooks
-        for raw_hook in plugin.raw_hooks:
+        for raw_hook in plugin.hooks["irc_raw"]:
             if raw_hook.is_catch_all():
                 self.catch_all_triggers.remove(raw_hook)
             else:
@@ -333,28 +327,32 @@ def unload_plugin(self, path):
                         del self.raw_triggers[trigger]
 
         # unregister events
-        for event_hook in plugin.events:
+        for event_hook in plugin.hooks["event"]:
             for event_type in event_hook.types:
                 assert event_type in self.event_type_hooks  # this can't be not true
                 self.event_type_hooks[event_type].remove(event_hook)
                 if not self.event_type_hooks[event_type]:  # if that was the last hook for this event type
                     del self.event_type_hooks[event_type]
 
         # unregister regexps
-        for regex_hook in plugin.regexes:
+        for regex_hook in plugin.hooks["regex"]:
             for regex_match in regex_hook.regexes:
                 self.regex_hooks.remove((regex_match, regex_hook))
 
         # unregister sieves
-        for sieve_hook in plugin.sieves:
+        for sieve_hook in plugin.hooks["sieve"]:
             self.sieves.remove(sieve_hook)
 
         # unregister connect hooks
-        for connect_hook in plugin.connect_hooks:
+        for connect_hook in plugin.hooks["on_connect"]:
             self.connect_hooks.remove(connect_hook)
 
+        for perm_hook in plugin.hooks["perm_check"]:
+            for perm in perm_hook.perms:
+                self.perm_hooks[perm].remove(perm_hook)
+
         # Run on_stop hooks
-        for on_stop_hook in plugin.run_on_stop:
+        for on_stop_hook in plugin.hooks["on_stop"]:
             event = Event(bot=self.bot, hook=on_stop_hook)
             yield from self.launch(on_stop_hook, event)
 
@@ -562,11 +560,7 @@ class Plugin:
     :type file_path: str
     :type file_name: str
     :type title: str
-    :type commands: list[CommandHook]
-    :type regexes: list[RegexHook]
-    :type raw_hooks: list[RawHook]
-    :type sieves: list[SieveHook]
-    :type events: list[EventHook]
+    :type hooks: dict
     :type tables: list[sqlalchemy.Table]
     """
 
@@ -580,12 +574,7 @@ def __init__(self, filepath, filename, title, code):
         self.file_path = filepath
         self.file_name = filename
         self.title = title
-        # 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
+        self.hooks = 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)
@@ -888,6 +877,18 @@ def __str__(self):
         return "{name} {func} from {file}".format(name=self.type, func=self.function_name, file=self.plugin.file_name)
 
 
+class PermHook(Hook):
+    def __init__(self, plugin, perm_hook):
+        self.perms = perm_hook.perms
+        super().__init__("perm_check", plugin, perm_hook)
+
+    def __repr__(self):
+        return "PermHook[{}]".format(Hook.__repr__(self))
+
+    def __str__(self):
+        return "perm hook {} from {}".format(self.function_name, self.plugin.file_name)
+
+
 _hook_name_to_plugin = {
     "command": CommandHook,
     "regex": RegexHook,
@@ -900,4 +901,5 @@ def __str__(self):
     "on_cap_available": OnCapAvaliableHook,
     "on_cap_ack": OnCapAckHook,
     "on_connect": OnConnectHook,
+    "perm_check": PermHook,
 }

cloudbot/util/async_util.py

@@ -5,6 +5,7 @@
 import asyncio
 
 import sys
+from functools import partial
 
 
 def wrap_future(fut, *, loop=None):
@@ -18,3 +19,12 @@ def wrap_future(fut, *, loop=None):
         return asyncio.async(fut, loop=loop)
 
     return asyncio.ensure_future(fut, loop=loop)
+
+
+@asyncio.coroutine
+def run_func(loop, func, *args, **kwargs):
+    part = partial(func, *args, **kwargs)
+    if asyncio.iscoroutine(func) or asyncio.iscoroutinefunction(func):
+        return (yield from part())
+    else:
+        return (yield from loop.run_in_executor(None, part))