rework how commands are processed, implement sudo commands

ubot/command_handler.py

@@ -6,35 +6,33 @@ from telethon import events
 
 
 class CommandHandler():
-    def __init__(self, client, logger):
+    def __init__(self, client, logger, settings):
+        self.pattern_template = "(?is)^{0}{1}(?: |$)(.*)"
         self.outgoing_commands = {}
         self.incoming_commands = {}
         self.logger = logger
-        client.add_event_handler(self.handle_outgoing, events.NewMessage(outgoing=True))
+        self.settings = settings
         client.add_event_handler(self.handle_incoming, events.NewMessage(incoming=True))
 
-    async def handle_outgoing(self, event):
-        for cmd in self.outgoing_commands.keys():
-            if search(cmd, event.text):
-                event.pattern_match = search(cmd, event.text)
-
-                try:
-                    await self.outgoing_commands.get(cmd)(event)
-                    return
-                except Exception as exception:
-                    self.logger.warn(f"{self.outgoing_commands.get(cmd).__name__} - {exception}")
-                    await event.reply(f"`An error occurred in {self.outgoing_commands.get(cmd).__name__}: {exception}`")
-                    raise exception
-
     async def handle_incoming(self, event):
-        for cmd in self.incoming_commands.keys():
-            if search(cmd, event.text):
-                event.pattern_match = search(cmd, event.text)
+        prefix = escape(self.settings.get_config("cmd_prefix") or '.')
+
+        for key, value in self.incoming_commands.items():
+            pattern = self.pattern_template.format("" if value["noprefix"] else prefix, key)
+
+            if search(pattern, event.text):
+                event.pattern_match = search(pattern, event.text)
+
+                if value["sudo"] and str(event.from_id) != self.settings.get_config("owner_id"):
+                    print(f"Attempted sudo command ({event.text}) from ID {event.from_id}")
+                    return
 
                 try:
-                    await self.incoming_commands.get(cmd)(event)
+                    await value["function"](event)
                     return
                 except Exception as exception:
-                    self.logger.warn(f"{self.incoming_commands.get(cmd).__name__} - {exception}")
-                    await event.reply(f"`An error occurred in {self.incoming_commands.get(cmd).__name__}: {exception}`")
+                    self.logger.warn(f"{value['function'].__name__} - {exception}")
+                    await event.reply(f"`An error occurred in {value['function'].__name__}: {exception}`")
                     raise exception
+
+                break

ubot/loader.py

@@ -17,7 +17,7 @@ class Loader():
         self.client = client
         self.logger = logger
         self.settings = settings
-        self.command_handler = CommandHandler(client, logger)
+        self.command_handler = CommandHandler(client, logger, settings)
         self.botversion = "0.1.3"
 
     def load_all_modules(self):
@@ -42,20 +42,12 @@ class Loader():
         return errors or None
 
     def add(self, **args):
-        prefix = escape(self.settings.get_config("cmd_prefix") or '.')
-
-        if args.get('noprefix', None):
-            del args['noprefix']
-            prefix = ''
-
-        if not args.get('isfilter', False) and args.get('pattern', None) is not None:
-            args['pattern'] = f"(?is)^{prefix}{args['pattern']}(?: |$|_)(.*)"
-        else:
-            del args['isfilter']
-            args['pattern'] = f"(?is)(.*){args['pattern']}(.*)"
-
         def decorator(func):
-            self.command_handler.incoming_commands[args['pattern']] = func
+            self.command_handler.incoming_commands[args['pattern']] = {
+                "function": func,
+                "noprefix": args.get('noprefix', False),
+                "sudo": args.get('sudo', False)
+            }
 
         return decorator
 

ubot/modules/system.py

@@ -10,22 +10,22 @@ from ubot.micro_bot import micro_bot
 ldr = micro_bot.loader
 
 
-#@ldr.add(pattern="reload")
+@ldr.add(pattern="reload", sudo=True)
 async def reload_modules(event):
-    await event.edit("`Reloading modules…`")
+    reload_msg = await event.reply("`Reloading modules…`")
 
     errors = ldr.reload_all_modules()
 
     if errors:
-        await event.edit(errors)
+        await reload_msg.edit(errors)
     else:
         try:
-            await event.delete()
+            await reload_msg.edit("`Successfully reloaded.`")
         except:
             pass
 
 
-@ldr.add(pattern="alive")
+@ldr.add(pattern="alive", sudo=True)
 async def alive(event):
     alive_format = "`μBot is running under {0}.\n\n" \
                    "Version: {1}\n" \
@@ -35,39 +35,12 @@ async def alive(event):
     await event.reply(alive_format.format(uname().node, ldr.botversion, version.__version__, python_version()))
 
 
-#@ldr.add(pattern="shutdown")
+@ldr.add(pattern="shutdown", sudo=True)
 async def shutdown(event):
-    await event.edit("`Goodbye…`")
+    await event.reply("`Goodbye…`")
     await micro_bot.stop_client()
 
 
-#@ldr.add(pattern="prefix")
-async def change_prefix(event):
-    new_prefix = event.pattern_match.group(1)
-
-    if not new_prefix:
-        await event.edit("`Please specify a valid command prefix!`")
-        return
-
-    micro_bot.settings.set_config("cmd_prefix", new_prefix)
-    errors = ldr.reload_all_modules()
-
-    if errors:
-        await event.edit(f"`Command prefix successfully changed to `**{new_prefix}**` but there were errors:`\n\n{errors}")
-    else:
-        await event.edit(f"`Command prefix successfully changed to `**{new_prefix}**`!`")
-
-
 @ldr.add(pattern="repo")
 async def bot_repo(event):
     await event.reply("https://github.com/Nick80835/microbot")
-
-
-#@ldr.add(pattern="toggleincoming")
-async def toggleincoming(event):
-    if micro_bot.settings.get_bool("incoming_allowed"):
-        micro_bot.settings.set_config("incoming_allowed", "False")
-        await event.edit("`Successfully disabled incoming commands!`")
-    else:
-        micro_bot.settings.set_config("incoming_allowed", "True")
-        await event.edit("`Successfully enabled incoming commands!`")