Queer European MD passionate about IT
Browse Source

Define roles and /restart command

Davte 5 years ago
parent
commit
2a455b6128
2 changed files with 88 additions and 1 deletions
  1. 25 1
      ciclopibot/bot.py
  2. 63 0
      ciclopibot/bot_tools.py

+ 25 - 1
ciclopibot/bot.py

@@ -3,14 +3,17 @@
 # Standard library modules
 import logging
 import os
+import sys
 
 # Third party modules
 from davtelepot.bot import Bot
 
 # Project modules
+import bot_tools
 import ciclopi
 from data.passwords import bot_token
 import helper
+import roles
 
 if __name__ == '__main__':
     path = os.path.dirname(__file__)
@@ -57,6 +60,25 @@ if __name__ == '__main__':
     # Instantiate bot
     bot = Bot(token=bot_token, database_url='ciclopibot/data/ciclopi.db')
     # Assign commands to bot
+    bot.set_unknown_command_message(
+        "Comando sconosciuto!\n"
+        "Scrivi /help per visualizzare la guida."
+    )
+    bot.set_authorization_function(
+        roles.get_authorization_function(bot)
+    )
+    bot.set_authorization_denied_message(
+        "Non disponi di autorizzazioni sufficienti per questo comando."
+    )
+    with bot.db as db:
+        db['users'].upsert(
+            dict(
+                telegram_id=63538990,
+                privileges=1
+            ),
+            ['telegram_id']
+        )
+    bot_tools.init(bot)
     ciclopi.init(bot)
     helper.init(
         bot=bot,
@@ -67,9 +89,11 @@ if __name__ == '__main__':
                      "Autore e amministratore del bot: @davte",
         help_sections_file='ciclopibot/data/help.json'
     )
+    roles.init(bot)
     # Run bot(s)
     logging.info("Presso ctrl+C to exit.")
-    Bot.run(
+    exit_state = Bot.run(
         local_host=local_host,
         port=port
     )
+    sys.exit(exit_state)

+ 63 - 0
ciclopibot/bot_tools.py

@@ -0,0 +1,63 @@
+"""Administration tools for CicloPiBot."""
+
+# Standard library modules
+import asyncio
+import datetime
+
+
+async def _restart_command(bot, update, user_record):
+    with bot.db as db:
+        db['restart_messages'].insert(
+            dict(
+                text="<i>Restart was successful.</i>",
+                chat_id=update['chat']['id'],
+                parse_mode='HTML',
+                reply_to_message_id=update['message_id'],
+                sent=None
+            )
+        )
+    await bot.reply(
+        update=update,
+        text="I bot verranno riavviati in pochi secondi, caricando prima le "
+             "eventuali modifiche al codice."
+    )
+    bot.__class__.stop(message='=== RESTART ===', final_state=65)
+    return
+
+
+def init(bot):
+    """Assign commands to `bot`."""
+    @bot.command(command='/restart', aliases=[], show_in_keyboard=False,
+                 description="Riavvia i bot",
+                 authorization_level='admin')
+    async def restart_command(bot, update, user_record):
+        return await _restart_command(bot, update, user_record)
+
+    @bot.additional_task('BEFORE')
+    async def load_handovers():
+        """Perform handovers before running."""
+        with bot.db as db:
+            for restart_message in db['restart_messages'].find(sent=None):
+                asyncio.ensure_future(
+                    bot.send_message(
+                        **{
+                            key: val
+                            for key, val in restart_message.items()
+                            if key in (
+                                'chat_id',
+                                'text',
+                                'parse_mode',
+                                'reply_to_message_id'
+                            )
+                        }
+                    )
+                )
+                db['restart_messages'].update(
+                    dict(
+                        sent=datetime.datetime.now(),
+                        id=restart_message['id']
+                    ),
+                    ['id'],
+                    ensure=True
+                )
+        return