From aff843c867311d3ffe852d38b32f4d1bf348c496 Mon Sep 17 00:00:00 2001 From: kiriharu Date: Tue, 1 Jun 2021 20:43:39 +0300 Subject: [PATCH 1/4] log warning if notification token not setted --- apps/tgbot/tgbot/handlers/helpers.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/apps/tgbot/tgbot/handlers/helpers.py b/apps/tgbot/tgbot/handlers/helpers.py index 48a8dbf..f2620cb 100644 --- a/apps/tgbot/tgbot/handlers/helpers.py +++ b/apps/tgbot/tgbot/handlers/helpers.py @@ -48,7 +48,10 @@ async def send_api_requests(endpoint: str, data: dict, nodes: List[APINode]): async def send_message_to_admins(message: str): - 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') + 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') + else: + logger.warning(f"Notificator bot token not setted. Skipping send notifications to admin") \ No newline at end of file From 3ca281627ff8fbe295b37de0b99784bd3a1bf70f Mon Sep 17 00:00:00 2001 From: kiriharu Date: Tue, 1 Jun 2021 21:01:33 +0300 Subject: [PATCH 2/4] Send all checks result to NOTIFICATION_USERS if NOTIFY_CHECKS enabled --- apps/tgbot/tgbot/config.py | 2 ++ apps/tgbot/tgbot/handlers/base.py | 9 ++++++++- apps/tgbot/tgbot/handlers/helpers.py | 2 +- 3 files changed, 11 insertions(+), 2 deletions(-) diff --git a/apps/tgbot/tgbot/config.py b/apps/tgbot/tgbot/config.py index 94b5b3d..aeb33e8 100644 --- a/apps/tgbot/tgbot/config.py +++ b/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 diff --git a/apps/tgbot/tgbot/handlers/base.py b/apps/tgbot/tgbot/handlers/base.py index 7baa880..dee57b8 100644 --- a/apps/tgbot/tgbot/handlers/base.py +++ b/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}"\ @@ -68,6 +69,12 @@ 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 {chat_id} issued check: {self.api_endpoint} for {data['target_fq']}**\n\n" \ + f"```{rsp_msg.text}```" + 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() diff --git a/apps/tgbot/tgbot/handlers/helpers.py b/apps/tgbot/tgbot/handlers/helpers.py index f2620cb..d316364 100644 --- a/apps/tgbot/tgbot/handlers/helpers.py +++ b/apps/tgbot/tgbot/handlers/helpers.py @@ -54,4 +54,4 @@ async def send_message_to_admins(message: str): logger.info(f"Sended notification to {user}") await bot.send_message(user, message, parse_mode='Markdown') else: - logger.warning(f"Notificator bot token not setted. Skipping send notifications to admin") \ No newline at end of file + logger.warning(f"Notificator bot token not setted. Skipping send notifications to admin") From 00ea38b251219e6255432a6fed6759ac51035070 Mon Sep 17 00:00:00 2001 From: kiriharu Date: Tue, 1 Jun 2021 21:16:51 +0300 Subject: [PATCH 3/4] small refactor --- apps/tgbot/tgbot/handlers/base.py | 13 ++++++++----- apps/tgbot/tgbot/handlers/default/icmp.py | 2 +- apps/tgbot/tgbot/handlers/default/start.py | 2 +- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/apps/tgbot/tgbot/handlers/base.py b/apps/tgbot/tgbot/handlers/base.py index dee57b8..75dc825 100644 --- a/apps/tgbot/tgbot/handlers/base.py +++ b/apps/tgbot/tgbot/handlers/base.py @@ -49,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( @@ -71,8 +74,9 @@ class CheckerBaseHandler(SimpleCommandHandler): iter_keys = iter_keys + 1 if NOTIFY_CHECKS: - notify_text = f"**User {chat_id} issued check: {self.api_endpoint} for {data['target_fq']}**\n\n" \ - f"```{rsp_msg.text}```" + 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}") @@ -108,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.bot, + message, dict(target=args[0], port=args[1], target_fq=f"{args[0]}:{args[1]}") ) diff --git a/apps/tgbot/tgbot/handlers/default/icmp.py b/apps/tgbot/tgbot/handlers/default/icmp.py index 6872832..4865219 100644 --- a/apps/tgbot/tgbot/handlers/default/icmp.py +++ b/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() diff --git a/apps/tgbot/tgbot/handlers/default/start.py b/apps/tgbot/tgbot/handlers/default/start.py index eaf267b..da867ed 100644 --- a/apps/tgbot/tgbot/handlers/default/start.py +++ b/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/) """ From 4ca6ae50cc2d102f78f3727394b29aed4f63e8b3 Mon Sep 17 00:00:00 2001 From: kiriharu Date: Tue, 1 Jun 2021 21:34:34 +0300 Subject: [PATCH 4/4] fix session closing bug --- apps/tgbot/tgbot/handlers/helpers.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/apps/tgbot/tgbot/handlers/helpers.py b/apps/tgbot/tgbot/handlers/helpers.py index d316364..04e89c0 100644 --- a/apps/tgbot/tgbot/handlers/helpers.py +++ b/apps/tgbot/tgbot/handlers/helpers.py @@ -47,11 +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")