implement update command and improved stop_client

requirements.txt

@@ -6,6 +6,7 @@ asyncpraw
 cchardet
 cryptg
 cython
+gitpython
 hachoir
 peewee
 pillow

ubot/__init__.py

@@ -29,11 +29,11 @@ class MicroBot():
         self.start_client()
         self.loader = Loader(self.client, self.logger, self.settings)
 
-    def run_until_done(self):
+    async def run_until_done(self):
         self.loader.load_all_modules()
         self.logger.info("Client successfully started.")
-        self.client.run_until_disconnected()
+        await self.client.run_until_disconnected()
-        self.client.loop.run_until_complete(self.loader.aioclient.close())
+        await self.stop_client(reason="Client disconnected.")
 
     def _check_config(self):
         api_key = self.settings.get_config("api_key")
@@ -64,16 +64,21 @@ class MicroBot():
             self.client = telethon.TelegramClient(self.settings.get_config("session_name", "bot0"), api_key, api_hash, connection=CTA).start(bot_token=bot_token)
         except (TokenInvalidError, AccessTokenExpiredError, AccessTokenInvalidError):
             self.logger.error("The bot token provided is invalid, exiting.")
-            sys.exit(2)
+            sys.exit(1)
 
-    async def stop_client(self, reason=None):
+    async def stop_client(self, reason=None, exit_code=None):
         if reason:
-            self.logger.info("Stopping client for reason: %s", reason)
+            self.logger.info("Stopping bot for reason: %s", reason)
         else:
-            self.logger.info("Stopping client.")
+            self.logger.info("Stopping bot.")
 
         await self.loader.aioclient.close()
-        await self.client.disconnect()
+
+        if exit_code is not None:
+            self.logger.info("Exiting with exit code: %i", exit_code)
+            sys.exit(exit_code)
+        else:
+            sys.exit(0)
 
 
 telethon.events.NewMessage.Event = ExtendedEvent
@@ -83,7 +88,4 @@ ldr = micro_bot.loader
 client = micro_bot.client
 logger = micro_bot.logger
 
-try:
+client.loop.run_until_complete(micro_bot.run_until_done())
-    micro_bot.run_until_done()
-except:
-    micro_bot.client.loop.run_until_complete(micro_bot.stop_client())

ubot/command_handler.py

@@ -5,7 +5,8 @@ from re import DOTALL, IGNORECASE, escape, search
 from traceback import format_exc, print_exc
 
 from telethon import events
-from telethon.errors.rpcerrorlist import ChatWriteForbiddenError, ChatAdminRequiredError
+from telethon.errors.rpcerrorlist import (ChatAdminRequiredError,
+                                          ChatWriteForbiddenError)
 
 from .database import ChatWrapper
 from .fixes import inline_photos

ubot/modules/_sudo.py

@@ -4,6 +4,7 @@ import io
 import os
 from platform import python_version
 
+import git
 import psutil
 from telethon import version
 from telethon.tl.types import Channel, Chat
@@ -89,6 +90,14 @@ async def reload_modules(event):
             pass
 
 
+@ldr.add("update", owner=True, hide_help=True)
+async def update_bot(event):
+    update_msg = await event.reply("Pulling changes…")
+    repo = git.Repo(os.getcwd())
+    repo.remotes.origin.pull()
+    await update_msg.edit("Changes pulled successfully!")
+
+
 @ldr.add("sysd", sudo=True, hide_help=True)
 async def sysd(event):
     try:
@@ -118,10 +127,14 @@ async def alive(event):
     await event.reply(alive_format.format(version.__version__, python_version(), mem_usage))
 
 
-@ldr.add("shutdown", owner=True, hide_help=True)
+@ldr.add("shutdown", pattern_extra="(f|)", owner=True, hide_help=True)
 async def shutdown(event):
     await event.reply("Goodbye…")
-    await micro_bot.stop_client()
+
+    if event.other_args[0]:
+        await micro_bot.stop_client(reason="Shutdown command issued.", exit_code=1)
+    else:
+        await micro_bot.stop_client(reason="Shutdown command issued.")
 
 
 @ldr.add("blacklist", sudo=True, hide_help=True)