From 352d27f8b2e2e2a4bae9621faa7ef2d338157cc8 Mon Sep 17 00:00:00 2001
From: kiriharu <kiriharu@yandex.ru>
Date: Wed, 13 Jan 2021 22:23:46 +0300
Subject: [PATCH] added CheckerTargetPortHandler, added validation in
 SimpleCommandHandler

---
 apps/tgbot/tgbot/handlers/base.py | 78 ++++++++++++++-----------------
 1 file changed, 36 insertions(+), 42 deletions(-)

diff --git a/apps/tgbot/tgbot/handlers/base.py b/apps/tgbot/tgbot/handlers/base.py
index f12a4d7..a87256f 100644
--- a/apps/tgbot/tgbot/handlers/base.py
+++ b/apps/tgbot/tgbot/handlers/base.py
@@ -1,12 +1,15 @@
 from aiogram.types import Message
-from typing import Tuple, Any
+from typing import Tuple, Any, List
 
 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 APINodeInfo
-from .helpers import send_api_requests, check_int, validate_local, timing
+from .helpers import send_api_requests, check_int
+from .errors import NotEnoughArgs, InvalidPort, LocalhostForbidden
+from .validators import BaseValidator, LocalhostValidator
+from tgbot.middlewares.throttling import rate_limit
 from loguru import logger
 from uuid import uuid4
 from time import time
@@ -16,21 +19,11 @@ header = "Отчет о проверке хоста:" \
          f"\n— Дата проверки: {datetime.now():%d.%m.%y в %H:%M} (MSK)"  # TODO: Get timezone
 
 
-class NotEnoughArgs(Exception):
-    pass
-
-
-class InvalidPort(Exception):
-    pass
-
-
-class LocalhostForbidden(Exception):
-    pass
-
-
 class SimpleCommandHandler:
     help_message = "Set help message in class!"
+    validators: List[BaseValidator] = [LocalhostValidator()]
 
+    @rate_limit
     async def handler(self, message: Message):
         pass
 
@@ -38,7 +31,8 @@ class SimpleCommandHandler:
         raise NotImplemented
 
     async def validate_target(self, target: str):
-        raise NotImplemented
+        for validator in self.validators:
+            validator.validate(target)
 
     async def prepare_message(self, *args) -> str:
         raise NotImplemented
@@ -53,28 +47,6 @@ class CheckerBaseHandler(SimpleCommandHandler):
     def __init__(self):
         pass
 
-    # TODO: create class CheckerTargetPortHandler
-    async def target_port_handler(self, message: Message):
-        """This hanler can be used if you need target port args"""
-        try:
-            args = await self.process_args(message.text)
-        except NotEnoughArgs:
-            logger.info(f"User {message.from_user.id} got NotEnoughArgs error")
-            return await message.answer(self.help_message, parse_mode="Markdown")
-        except InvalidPort:
-            logger.info(f"User {message.from_user.id} got InvalidPort error")
-            return await message.answer(self.invalid_port_message, parse_mode="Markdown")
-        try:
-            await self.validate_target(args[0])
-        except LocalhostForbidden:
-            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,
-            dict(target=args[0], port=args[1], target_fq=f"{args[0]}:{args[1]}")
-        )
-
     async def check(self, chat_id: int, bot: Bot, data: dict):
         # TODO: start check and end check metrics with ident, chat_id and api_endpoint
         ts = time()
@@ -96,11 +68,6 @@ class CheckerBaseHandler(SimpleCommandHandler):
         te = time()
         logger.info(f"func {__name__} took {te - ts} sec")
 
-    # TODO: Validation classes
-    async def validate_target(self, target: str):
-        if validate_local(target):
-            raise LocalhostForbidden()
-
     async def message_std_vals(self, res: Response) -> Tuple[str, Any]:
         node = APINodeInfo(**res.json().get("node", None))
         message = f"{node.location}:\n"
@@ -108,6 +75,33 @@ class CheckerBaseHandler(SimpleCommandHandler):
         return message, status
 
 
+class CheckerTargetPortHandler(CheckerBaseHandler):
+
+    @rate_limit
+    async def handler(self, message: Message):
+        """This hanler can be used if you need target port args"""
+        try:
+            args = await self.process_args(message.text)
+        except NotEnoughArgs:
+            logger.info(f"User {message.from_user.id} got NotEnoughArgs error")
+            return await message.answer(self.help_message, parse_mode="Markdown")
+        except InvalidPort:
+            logger.info(f"User {message.from_user.id} got InvalidPort error")
+            return await message.answer(self.invalid_port_message, parse_mode="Markdown")
+        try:
+            await self.validate_target(args[0])
+        except ValueError:  # For ip check
+            pass
+        except LocalhostForbidden:
+            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,
+            dict(target=args[0], port=args[1], target_fq=f"{args[0]}:{args[1]}")
+        )
+
+
 def process_args_for_host_port(text: str, default_port: int) -> list:
     port = None
     args = text.split()