From 3f39e88394eea76037ee3a89be731101ce3337a3 Mon Sep 17 00:00:00 2001
From: Dan Elkouby <streetwalkermc@gmail.com>
Date: Sun, 8 Apr 2018 01:09:02 +0300
Subject: [PATCH] Add core plugin with basic admin commands

---
 uniborg/_core.py    | 39 +++++++++++++++++++++++++++++++++++++++
 uniborg/telethon.py | 13 +++++++++++++
 uniborg/util.py     | 10 ++++++++++
 3 files changed, 62 insertions(+)
 create mode 100644 uniborg/util.py

diff --git a/uniborg/_core.py b/uniborg/_core.py
index 5501cd4..4c726ab 100644
--- a/uniborg/_core.py
+++ b/uniborg/_core.py
@@ -2,3 +2,42 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
+import asyncio
+import traceback
+
+from uniborg import util
+
+DELETE_TIMEOUT = 2
+
+@borg.on(util.admin_cmd(r"^\.load (?P<shortname>\w+)$"))
+async def load_reload(event):
+    await event.delete()
+    shortname = event.pattern_match["shortname"]
+
+    try:
+        if shortname in borg._plugins:
+            borg.remove_plugin(shortname)
+        borg.load_plugin(shortname)
+
+        msg = await event.respond(f"Successfully (re)loaded plugin {shortname}")
+        await asyncio.sleep(DELETE_TIMEOUT)
+        await borg.delete_messages(msg.to_id, msg)
+
+    except Exception as e:
+        tb = traceback.format_exc()
+        logger.warn(f"Failed to (re)load plugin {shortname}: {tb}")
+        await event.respond(f"Failed to (re)load plugin {shortname}: {e}")
+
+@borg.on(util.admin_cmd(r"^\.remove (?P<shortname>\w+)$"))
+async def remove(event):
+    await event.delete()
+    shortname = event.pattern_match["shortname"]
+
+    if shortname in borg._plugins:
+        borg.remove_plugin(shortname)
+        msg = await event.respond(f"Removed plugin {shortname}")
+    else:
+        msg = await event.respond(f"Plugin {shortname} is not loaded")
+
+    await asyncio.sleep(DELETE_TIMEOUT)
+    await borg.delete_messages(msg.to_id, msg)
diff --git a/uniborg/telethon.py b/uniborg/telethon.py
index b258338..b31080a 100644
--- a/uniborg/telethon.py
+++ b/uniborg/telethon.py
@@ -62,3 +62,16 @@ class Uniborg(TelegramClient):
         spec.loader.exec_module(mod)
         self._plugins[shortname] = mod
         self._logger.info(f"Successfully loaded plugin {shortname}")
+
+    def remove_plugin(self, shortname):
+        name = self._plugins[shortname].__name__
+
+        i = len(self._event_builders)
+        while i:
+            i -= 1
+            ev, cb = self._event_builders[i]
+            if cb.__module__ == name:
+                del self._event_builders[i]
+
+        del self._plugins[shortname]
+        self._logger.info(f"Removed plugin {shortname}")
diff --git a/uniborg/util.py b/uniborg/util.py
new file mode 100644
index 0000000..0056e8e
--- /dev/null
+++ b/uniborg/util.py
@@ -0,0 +1,10 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+import re
+
+from telethon import events
+
+def admin_cmd(pattern):
+    return events.NewMessage(outgoing=True, pattern=re.compile(pattern))