feat: add maintenance mode commands

milselarch · Nov 22, 2024 · c621c3f · c621c3f
1 parent 4dab03c
commit c621c3f
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 9 deletions.
diff --git a/bot.py b/bot.py
@@ -84,6 +84,8 @@ def __init__(self, config_path='config.yml'):
         super().__init__()
         self.config_path = config_path
         self.scheduled_processes = []
+        self.payment_handlers = None
+
         self.webhook_url = None
         self.bot = None
         self.app = None
@@ -122,7 +124,7 @@ def start_bot(self):
         builder.post_init(self.post_init)
 
         self.app = builder.build()
-        payment_handlers = PaymentHandlers(logger)
+        self.payment_handlers = PaymentHandlers(logger)
 
         commands_mapping = {
             Command.START: start_handlers.start_handler,
@@ -154,17 +156,19 @@ def start_bot(self):
             Command.LOOKUP_FROM_USERNAME_ADMIN:
                 self.lookup_from_username_admin,
             Command.INSERT_USER_ADMIN: self.insert_user_admin,
-            Command.SET_MAX_VOTERS: payment_handlers.set_max_voters,
+            Command.SET_MAX_VOTERS: self.payment_handlers.set_max_voters,
             Command.PAY_SUPPORT: self.payment_support_handler,
-            Command.REFUND_ADMIN: self.refund_payment_support_handler
+            Command.REFUND_ADMIN: self.refund_payment_support_handler,
+            Command.ENTER_MAINTENANCE_ADMIN: self.enter_maintenance_admin,
+            Command.EXIT_MAINTENANCE_ADMIN: self.exit_maintenance_admin
         }
 
         # on different commands - answer in Telegram
         TelegramHelpers.register_commands(
             self.app, commands_mapping=commands_mapping
         )
         TelegramHelpers.register_pre_checkout_handler(
-            self.app, payment_handlers.pre_checkout_callback
+            self.app, self.payment_handlers.pre_checkout_callback
         )
         # catch-all to handle responses to unknown commands
         TelegramHelpers.register_message_handler(
@@ -179,7 +183,7 @@ def start_bot(self):
         # handle payment-related messages
         TelegramHelpers.register_message_handler(
             self.app, filters.SUCCESSFUL_PAYMENT,
-            payment_handlers.successful_payment_callback
+            self.payment_handlers.successful_payment_callback
         )
         # catch-all to handle all other messages
         TelegramHelpers.register_message_handler(
@@ -1670,6 +1674,22 @@ async def refund_payment_support_handler(
             payment.refunded_at = datetime.now()
             payment.save()
 
+    @admin_only
+    def enter_maintenance_admin(
+        self, update: ModifiedTeleUpdate, _: ContextTypes.DEFAULT_TYPE
+    ):
+        message = update.message
+        self.payment_handlers.enter_maintenance_mode()
+        return message.reply_text('Maintenance mode entered')
+
+    @admin_only
+    def exit_maintenance_admin(
+        self, update: ModifiedTeleUpdate, _: ContextTypes.DEFAULT_TYPE
+    ):
+        message = update.message
+        self.payment_handlers.exit_maintenance_mode()
+        return message.reply_text('Maintenance mode exited')
+
 
 if __name__ == '__main__':
     rcv_bot = RankedChoiceBot()

diff --git a/handlers/payment_handlers.py b/handlers/payment_handlers.py
@@ -204,10 +204,17 @@ async def successful_payment_callback(
 class PaymentHandlers(object):
     def __init__(self, logger: logging.Logger):
         self.logger = logger
+        self.maintenance = False
         self.handlers: dict[InvoiceTypes, Type[BasePaymentHandler]] = {
             InvoiceTypes.INCREASE_VOTER_LIMIT: IncreaseVoteLimitHandler
         }
 
+    def enter_maintenance_mode(self):
+        self.maintenance = True
+
+    def exit_maintenance_mode(self):
+        self.maintenance = False
+
     async def successful_payment_callback(
         self, update: ModifiedTeleUpdate, context: ContextTypes.DEFAULT_TYPE
     ):
@@ -257,16 +264,18 @@ async def successful_payment_callback(
     async def pre_checkout_callback(
         self, update: ModifiedTeleUpdate, context: ContextTypes.DEFAULT_TYPE
     ):
-        # TODO: validate receipt has not expired
         query = update.pre_checkout_query
+        async def fail(err_message: str):
+            return await query.answer(ok=False, error_message=err_message)
+
+        if self.maintenance:
+            return await fail("Bot is in maintenance mode")
+
         invoice_payload = query.invoice_payload
         base_invoice_params_res = BasePaymentParams.safe_load_from_json(
             invoice_payload
         )
 
-        async def fail(err_message: str):
-            return await query.answer(ok=False, error_message=err_message)
-
         if base_invoice_params_res.is_err():
             self.logger.error(f"LOAD INVOICE FAILED: {invoice_payload}")
             return await fail("Error loading invoice info")

diff --git a/helpers/commands.py b/helpers/commands.py
@@ -35,3 +35,5 @@ class Command(StrEnum):
     LOOKUP_FROM_USERNAME_ADMIN = "lookup_from_username_admin"
     INSERT_USER_ADMIN = "insert_user_admin"
     REFUND_ADMIN = "refund_admin"
+    ENTER_MAINTENANCE_ADMIN = "enter_maintenance_admin"
+    EXIT_MAINTENANCE_ADMIN = "exit_maintenance_admin"