From 09fd459a1f50936c48adb1ef02370dda65111f5a Mon Sep 17 00:00:00 2001
From: kiriharu <kiriharu@yandex.ru>
Date: Wed, 13 Jan 2021 22:36:05 +0300
Subject: [PATCH] starting rewriting whois

---
 apps/tgbot/tgbot/handlers/base.py             |  2 +-
 apps/tgbot/tgbot/handlers/default/__init__.py |  4 +-
 apps/tgbot/tgbot/handlers/default/whois.py    | 46 +++++++++++++------
 3 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/apps/tgbot/tgbot/handlers/base.py b/apps/tgbot/tgbot/handlers/base.py
index a87256f..623e850 100644
--- a/apps/tgbot/tgbot/handlers/base.py
+++ b/apps/tgbot/tgbot/handlers/base.py
@@ -21,6 +21,7 @@ header = "Отчет о проверке хоста:" \
 
 class SimpleCommandHandler:
     help_message = "Set help message in class!"
+    localhost_forbidden_message = "❗ Локальные адреса запрещены"
     validators: List[BaseValidator] = [LocalhostValidator()]
 
     @rate_limit
@@ -39,7 +40,6 @@ class SimpleCommandHandler:
 
 
 class CheckerBaseHandler(SimpleCommandHandler):
-    localhost_forbidden_message = "❗ Локальные адреса запрещены"
     invalid_port_message = "Invalid port!"
     header_message = header
     api_endpoint = "Set api endpoint in class!"
diff --git a/apps/tgbot/tgbot/handlers/default/__init__.py b/apps/tgbot/tgbot/handlers/default/__init__.py
index 97d125a..bafa9fc 100644
--- a/apps/tgbot/tgbot/handlers/default/__init__.py
+++ b/apps/tgbot/tgbot/handlers/default/__init__.py
@@ -2,7 +2,7 @@ from aiogram import Dispatcher
 
 from .start import start_cmd
 from .web import WebCheckerHandler
-from .whois import whois_cmd
+from .whois import WhoisCommandHandler
 from .icmp import ICMPCheckerHandler
 from .tcp import TCPCheckerHandler
 from .minecraft import MinecraftCheckerHandler
@@ -12,7 +12,7 @@ from .ipcalc import IPCalcCommandHandler
 def setup(dp: Dispatcher):
     dp.register_message_handler(start_cmd, is_forwarded=False, commands=['start'])
     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(WhoisCommandHandler().handler, is_forwarded=False, commands=['whois'])
     dp.register_message_handler(ICMPCheckerHandler().handler, is_forwarded=False, commands=['icmp', 'ping'])
     dp.register_message_handler(TCPCheckerHandler().handler, is_forwarded=False, commands=['tcp'])
     dp.register_message_handler(MinecraftCheckerHandler().handler, is_forwarded=False, commands=['minecraft'])
diff --git a/apps/tgbot/tgbot/handlers/default/whois.py b/apps/tgbot/tgbot/handlers/default/whois.py
index d15d716..5b93991 100644
--- a/apps/tgbot/tgbot/handlers/default/whois.py
+++ b/apps/tgbot/tgbot/handlers/default/whois.py
@@ -2,7 +2,8 @@ from whois import whois, parser
 from aiogram.types import Message
 from aiogram.utils.markdown import quote_html
 
-from tgbot.handlers.helpers import validate_local
+from tgbot.handlers.base import SimpleCommandHandler
+from tgbot.handlers.errors import NotEnoughArgs, LocalhostForbidden
 from tgbot.middlewares.throttling import rate_limit
 
 whois_help_message = """
@@ -17,10 +18,7 @@ no_domain_text = """
 Напишите /whois чтобы посмотреть справку.
 """
 
-localhost_exception = "❗Локальные адреса запрещены!"
 
-
-# TODO: Very shitty code. I should rewrite this.
 def create_whois_message(domain: str) -> str:
     try:
         domain_info = whois(domain)
@@ -80,14 +78,32 @@ def create_whois_message(domain: str) -> str:
     return message
 
 
-@rate_limit
-async def whois_cmd(msg: Message):
-    args = msg.text.split()
-    if len(args) == 1:
-        return await msg.answer(no_domain_text)
-    if len(args) >= 2:
-        host = args[1]
-        if validate_local(args[0]):
-            return await msg.answer(localhost_exception, parse_mode="Markdown")
-        await msg.bot.send_chat_action(msg.chat.id, 'typing')
-        await msg.answer(create_whois_message(host), parse_mode='html')
+class WhoisCommandHandler(SimpleCommandHandler):
+
+    help_message = whois_help_message
+
+    def __init__(self):
+        super().__init__()
+
+    @rate_limit
+    async def handler(self, message: Message):
+        try:
+            args = await self.process_args(message.text)
+        except NotEnoughArgs:
+            await message.answer(no_domain_text, parse_mode='Markdown')
+        except LocalhostForbidden:
+            await message.answer(self.localhost_forbidden_message, parse_mode='Markdown')
+        else:
+            await message.answer(create_whois_message(args[0]), parse_mode='html')
+
+    async def process_args(self, text: str) -> list:
+        args = text.split()
+        if len(args) == 1:
+            raise NotEnoughArgs
+        if len(args) >= 2:
+            host = args[1]
+            await self.validate_target(host)
+            return [host]  # only domain name
+
+    async def prepare_message(self) -> str:
+        pass