impl #2, added tcp method

apps/tgbot/tgbot/handlers/base.py

@@ -5,11 +5,11 @@ from tgbot.nodes import nodes as all_nodes
 from httpx import Response
 from aiogram.bot import Bot
 from datetime import datetime
-from core.coretypes import ErrorPayload, ICMPCheckerResponse, ResponseStatus, APINodeInfo
+from core.coretypes import APINodeInfo
 from .helpers import send_api_requests
 
 header = "Отчет о проверке хоста:" \
-         "\n\n— Хост: {target}"\
+         "\n\n— Хост: {target_fq}"\
          f"\n— Дата проверки: {datetime.now():%d.%m.%y в %H:%M} (MSK)"  # TODO: Get timezone
 
 
@@ -30,6 +30,7 @@ class CheckerBaseHandler:
         pass
 
     async def handler(self, message: Message):
+        """Always should call check at end"""
         raise NotImplemented
 
     async def check(self, chat_id: int, bot: Bot, data: dict):
@@ -39,7 +40,7 @@ class CheckerBaseHandler:
         async for res in send_api_requests(self.api_endpoint, data, all_nodes):
             await bot.send_chat_action(chat_id, 'typing')
             if res.status_code == 500:
-                rsp_msg = await rsp_msg.edit_text(rsp_msg.text + f"\n\n{iter_keys}. Backend offline!")
+                rsp_msg = await rsp_msg.edit_text(rsp_msg.text + f"\n\n{iter_keys}. ❌️ Результат операции не доступен.")
             else:
                 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}")

apps/tgbot/tgbot/handlers/default/__init__.py

@@ -1,13 +1,15 @@
 from aiogram import Dispatcher
 
 from .start import start_cmd
-from .web import web_cmd
+from .web import WebCheckerHandler
 from .whois import whois_cmd
-from .icmp import icmp_cmd
+from .icmp import ICMPCheckerHandler
+from .tcp import TCPCheckerHandler
 
 
 def setup(dp: Dispatcher):
     dp.register_message_handler(start_cmd, is_forwarded=False, commands=['start'])
-    dp.register_message_handler(web_cmd, is_forwarded=False, commands=['web', 'http'])
+    dp.register_message_handler(WebCheckerHandler().handler, is_forwarded=False, commands=['web', 'http'])
     dp.register_message_handler(whois_cmd, is_forwarded=False, commands=['whois'])
-    dp.register_message_handler(icmp_cmd, is_forwarded=False, commands=['icmp', 'ping'])
+    dp.register_message_handler(ICMPCheckerHandler().handler, is_forwarded=False, commands=['icmp', 'ping'])
+    dp.register_message_handler(TCPCheckerHandler().handler, is_forwarded=False, commands=['tcp'])

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

@@ -1,9 +1,7 @@
 from aiogram.types import Message
-from tgbot.nodes import nodes as all_nodes
+from httpx import Response
-from httpx import AsyncClient, Response
+from core.coretypes import ErrorPayload, ICMPCheckerResponse, ResponseStatus
-from datetime import datetime
+from ..base import CheckerBaseHandler, NotEnoughArgs
-from core.coretypes import ErrorCodes, ErrorPayload, ICMPCheckerResponse, ResponseStatus, APINodeInfo
-from ..helpers import send_api_requests
 
 icmp_help_message = """
 ❓ Производит проверку хоста по протоколу ICMP.
@@ -13,11 +11,30 @@ icmp_help_message = """
 """
 
 
-async def prepare_icmp_check_result(res: Response):
+class ICMPCheckerHandler(CheckerBaseHandler):
-    node = APINodeInfo(**res.json().get("node", None))
+    help_message = icmp_help_message
-    message = f"{node.location}:\n"
+    api_endpoint = "/icmp"
-    status = res.json().get("status", None)
 
+    def __init__(self):
+        super(ICMPCheckerHandler, self).__init__()
+
+    async def handler(self, message: Message):
+        try:
+            args = await self.process_args(message.text)
+        except NotEnoughArgs:
+            return await message.answer(icmp_help_message)
+        await self.check(message.chat.id, message.bot, dict(target=args[0], target_fq=args[0]))
+
+    async def process_args(self, text: str) -> list:
+        args = text.split(" ")
+        if len(args) == 1:
+            raise NotEnoughArgs()
+        if len(args) >= 2:
+            target = args[1]
+            return [target]
+
+    async def prepare_message(self, res: Response):
+        message, status = await self.message_std_vals(res)
         if status == ResponseStatus.OK:
             payload = ICMPCheckerResponse(**res.json().get("payload"))
             message += f"✅ {payload.min_rtt}/{payload.max_rtt}/{payload.avg_rtt} " \
@@ -26,27 +43,3 @@ async def prepare_icmp_check_result(res: Response):
             payload = ErrorPayload(**res.json().get("payload"))
             message += f"❌️ {payload.message}"
         return message
-
-
-async def check_icmp(msg: Message, target: str):
-    rsp_msg = await msg.answer(f"Отчет о проверке хоста:"
-                               f"\n\n— Хост: {target}"
-                               f"\n— Дата проверки: {datetime.now():%d.%m.%y в %H:%M} (MSK)"  # TODO: Get timezone
-                               )
-    iter_keys = 1  # because I can't use enumerate
-    # using generators for magic...
-    async for res in send_api_requests("icmp", dict(target=target), all_nodes):
-        await msg.bot.send_chat_action(msg.chat.id, 'typing')
-        node_formatted_response = await prepare_icmp_check_result(res)
-        rsp_msg = await rsp_msg.edit_text(rsp_msg.text + f"\n\n{iter_keys}. {node_formatted_response}")
-        iter_keys = iter_keys + 1
-    await rsp_msg.edit_text(rsp_msg.text + f"\n\nПроверка завершена❗")
-
-
-async def icmp_cmd(msg: Message):
-    args = msg.text.split(" ")
-    if len(args) == 1:
-        return await msg.answer(icmp_help_message, parse_mode="Markdown")
-    if len(args) >= 2:
-        target = args[1]
-        await check_icmp(msg, target)

apps/tgbot/tgbot/handlers/default/tcp.py

@@ -0,0 +1,61 @@
+from aiogram.types import Message
+from core.coretypes import ResponseStatus, ErrorPayload, PortResponse
+from httpx import Response
+
+from tgbot.handlers.base import CheckerBaseHandler, NotEnoughArgs, InvalidPort
+from tgbot.handlers.helpers import check_int
+
+tcp_help_message = """
+❓ Производит проверку TCP порта, открыт ли он или нет
+
+Использование:
+ `/tcp <hostname> <port>` 
+"""
+
+invalid_port = """❗Неправильный порт. Напишите /tcp чтобы увидеть справку к данному способу проверки."""
+
+
+class TCPCheckerHandler(CheckerBaseHandler):
+    help_message = tcp_help_message
+    api_endpoint = "/tcp_port"
+
+    def __init__(self):
+        super().__init__()
+
+    async def handler(self, message: Message):
+        try:
+            args = await self.process_args(message.text)
+        except NotEnoughArgs:
+            return await message.answer(self.help_message, parse_mode="Markdown")
+        except InvalidPort:
+            return await message.answer(invalid_port, parse_mode="Markdown")
+        await self.check(
+            message.chat.id,
+            message.bot,
+            dict(target=args[0], port=args[1], target_fq=f"{args[0]}:{args[1]}")
+        )
+
+    async def process_args(self, text: str) -> list:
+        port = None
+        args = text.split(" ")
+        if len(args) < 3:
+            raise NotEnoughArgs()
+        if len(args) >= 3:
+            port = args[2]
+            if not check_int(port):
+                raise InvalidPort()
+        host = args[1]
+        return [host, port]
+
+    async def prepare_message(self, res: Response):
+        message, status = await self.message_std_vals(res)
+        if status == ResponseStatus.OK:
+            payload = PortResponse(**res.json().get("payload"))
+            if payload.open:
+                message += "✅ Порт ОТКРЫТ"
+            else:
+                message += "❌️ Порт ЗАКРЫТ"
+        if status == ResponseStatus.ERROR:
+            payload = ErrorPayload(**res.json().get("payload"))
+            message += f"❌️ {payload.message}"
+        return message

apps/tgbot/tgbot/handlers/default/web.py

@@ -1,11 +1,8 @@
 from aiogram.types import Message
-from typing import Optional
 from tgbot.handlers.helpers import check_int
-from tgbot.nodes import nodes as all_nodes
 from httpx import Response
-from core.coretypes import ResponseStatus, HTTP_EMOJI
+from core.coretypes import ResponseStatus, HTTP_EMOJI, HttpCheckerResponse, ErrorPayload
-from datetime import datetime
+from ..base import CheckerBaseHandler, NotEnoughArgs, InvalidPort
-from ..helpers import send_api_requests
 
 web_help_message = """
 ❓ Производит проверку хоста по протоколу HTTP.
@@ -18,55 +15,47 @@ web_help_message = """
 invalid_port = """❗Неправильный порт. Напишите /web чтобы увидеть справку к данному способу проверки."""
 
 
-async def prepare_webcheck_message(response: Response) -> str:
+class WebCheckerHandler(CheckerBaseHandler):
-    # TODO: Use types from core!
+    help_message = web_help_message
-    message = ""
+    api_endpoint = "/http"
-    json_rsp = response.json()
-    status = json_rsp.get("status")
-    location = json_rsp['node']['location']
-    if status == ResponseStatus.OK:
-        status_code = json_rsp['payload']['status_code']
-        time = round(json_rsp['payload']['time'], 2)
-        message = f"{location}:" \
-                  f"\n{HTTP_EMOJI.get(status_code//100, '')} {status_code}, ⏰ {time} сек."
-    if status == ResponseStatus.ERROR:
-        message = json_rsp['payload']['message']
-        message = f"{location}: " \
-                  f"\n❌ {message}"
-    return message
 
+    def __init__(self):
+        super().__init__()
 
-async def check_web(message: Message, host: str, port: Optional[int]):
+    async def handler(self, message: Message):
-    if port is None:
+        try:
-        port = 80
+            args = await self.process_args(message.text)
-    rsp_msg = await message.answer(f"Отчет о проверке хоста:"
+        except NotEnoughArgs:
-                                   f"\n\n— Хост: {host}:{port}"
+            return await message.answer(self.help_message, parse_mode="Markdown")
-                                   f"\n— Дата проверки: {datetime.now():%d.%m.%y в %H:%M} (MSK)"  # TODO: Get timezone
+        except InvalidPort:
+            return await message.answer(invalid_port, parse_mode="Markdown")
+        await self.check(
+            message.chat.id,
+            message.bot,
+            dict(target=args[0], port=args[1], target_fq=f"{args[0]}:{args[1]}")
         )
-    iter_keys = 1  # because I can't use enumerate
-    # using generators for magic...
-    async for res in send_api_requests("http", dict(target=host, port=port), all_nodes):
-        # set typing status...
-        await message.bot.send_chat_action(message.chat.id, 'typing')
-
-        node_formatted_response = await prepare_webcheck_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
-    await rsp_msg.edit_text(rsp_msg.text + f"\n\nПроверка завершена❗")
-
-
-async def web_cmd(msg: Message):
 
+    async def process_args(self, text: str) -> list:
         port = None
-    # TODO: Maybe check it in separated function?
+        args = text.split(" ")
-    args = msg.text.split(" ")
         if len(args) < 2:
-        return await msg.answer(web_help_message, parse_mode="Markdown")
+            raise NotEnoughArgs()
         if len(args) == 3:
             port = args[2]
             if not check_int(port):
-            return await msg.answer(invalid_port, parse_mode="Markdown")
+                raise InvalidPort()
+        if len(args) == 2:
+            port = 80
         host = args[1]
+        return [host, port]
 
-    await check_web(msg, host, port)
+    async def prepare_message(self, res: Response):
-
+        message, status = await self.message_std_vals(res)
+        if status == ResponseStatus.OK:
+            payload = HttpCheckerResponse(**res.json().get("payload"))
+            message += f"{HTTP_EMOJI.get(payload.status_code // 100, '')} " \
+                       f"{payload.status_code}, ⏰ {payload.time * 1000:.2f}ms"
+        if status == ResponseStatus.ERROR:
+            payload = ErrorPayload(**res.json().get("payload"))
+            message += f"❌️ {payload.message}"
+        return message

apps/tgbot/tgbot/handlers/helpers.py

@@ -21,6 +21,7 @@ async def send_api_requests(endpoint: str, data: dict, nodes: List[APINode]):
                     f"{node.address}/{endpoint}", params=data
                 )
         except ConnectError:
+            # TODO: Report problems to admins
             # We yield 500 response when backend is offline
             result = Response(500)
         yield result