Merge pull request #34 from catspace-dev/notificators

small refactor, notify checks to admins

apps/tgbot/tgbot/config.py

@@ -13,6 +13,8 @@ INFLUX_DB = os.getenv("INFLUX_DB", None)
 # Notifications
 NOTIFICATION_BOT_TOKEN = os.getenv("NOTIFICATION_BOT_TOKEN")
 NOTIFICATION_USERS = os.getenv("NOTIFICATION_USERS", "").split(",")
+# Send all checks result to NOTIFICATION_USERS
+NOTIFY_CHECKS = True
 
 # Mysql params
 MYSQL_HOST = os.getenv("MYSQL_HOST", None)  # if none, use sqlite db

apps/tgbot/tgbot/handlers/base.py

@@ -12,8 +12,9 @@ from loguru import logger
 from ..middlewares.throttling import rate_limit
 from ..nodes import nodes as all_nodes
 from .errors import InvalidPort, LocalhostForbidden, NotEnoughArgs
-from .helpers import send_api_requests
+from .helpers import send_api_requests, send_message_to_admins
 from .validators import BaseValidator, LocalhostValidator
+from ..config import NOTIFY_CHECKS
 
 header = "Отчет о проверке хоста:" \
          "\n\n— Хост: {0}"\
@@ -48,10 +49,13 @@ class CheckerBaseHandler(SimpleCommandHandler):
     def __init__(self):
         pass
 
-    async def check(self, chat_id: int, bot: Bot, data: dict):
+    async def check(self, msg: Message, data: dict):
         # TODO: start check and end check metrics with ident, chat_id and api_endpoint
         ts = time()
         ident = uuid4().hex
+        # refactoring goes brr
+        chat_id = msg.chat.id
+        bot = msg.bot
         logger.info(f"User {chat_id} started check {ident}")
         # format header
         rsp_msg = await bot.send_message(
@@ -68,6 +72,13 @@ class CheckerBaseHandler(SimpleCommandHandler):
                 node_formatted_response = await self.prepare_message(res)
                 rsp_msg = await rsp_msg.edit_text(rsp_msg.text + f"\n\n{iter_keys}. {node_formatted_response}")
             iter_keys = iter_keys + 1
+
+        if NOTIFY_CHECKS:
+            notify_text = f"**User {msg.from_user.full_name} (@{msg.from_user.username}) ({chat_id}) issued check: " \
+                          f"{self.api_endpoint} for {data['target_fq']}**\n\n" \
+                   f"```\n{rsp_msg.text}\n```"
+            await send_message_to_admins(notify_text)
+
         logger.info(f"User {chat_id} ended check {ident}")
         await rsp_msg.edit_text(rsp_msg.text + f"\n\nПроверка завершена❗")
         te = time()
@@ -101,8 +112,7 @@ class CheckerTargetPortHandler(CheckerBaseHandler):
             logger.info(f"User {message.from_user.id} got LocalhostForbidden error")
             return await message.answer(self.localhost_forbidden_message, parse_mode="Markdown")
         await self.check(
-            message.chat.id,
+            message,
-            message.bot,
             dict(target=args[0], port=args[1], target_fq=f"{args[0]}:{args[1]}")
         )
 

apps/tgbot/tgbot/handlers/default/icmp.py

@@ -30,7 +30,7 @@ class ICMPCheckerHandler(CheckerBaseHandler):
         except LocalhostForbidden:
             return await message.answer(self.localhost_forbidden_message, parse_mode="Markdown")
         else:
-            await self.check(message.chat.id, message.bot, dict(target=args[0], target_fq=args[0]))
+            await self.check(message, dict(target=args[0], target_fq=args[0]))
 
     def process_args(self, text: str) -> list:
         args = text.split()

apps/tgbot/tgbot/handlers/default/start.py

@@ -27,7 +27,7 @@ start_message = """
 🚩 [Этот бот с открытым с исходным кодом](https://github.com/catspace-dev/unicheckbot)
 🚩 [Помогите улучшить бота](https://github.com/catspace-dev/unicheckbot/issues) или [расскажите об ошибке](https://github.com/catspace-dev/unicheckbot/issues)
 
-Разработчик: [kiriharu](http://t.me/kiriharu)
+Разработчик: [kiriharu](https://t.me/kiriharu)
 При поддержке: [SpaceCore.pro](https://spacecore.pro/)
 
 """

apps/tgbot/tgbot/handlers/helpers.py

@@ -47,8 +47,12 @@ async def send_api_requests(endpoint: str, data: dict, nodes: List[APINode]):
             yield res
 
 
-async def send_message_to_admins(message: str):
+async def send_message_to_admins(message: str) -> None:
+    if NOTIFICATION_BOT_TOKEN:
         bot = Bot(token=NOTIFICATION_BOT_TOKEN)
         for user in NOTIFICATION_USERS:
             logger.info(f"Sended notification to {user}")
             await bot.send_message(user, message, parse_mode='Markdown')
+        await bot.close()
+    else:
+        logger.warning(f"Notificator bot token not setted. Skipping send notifications to admin")