Merge pull request #29 from cmd410/imports_fix

This should probably fix relative imports in tgbot

apps/Dockerfile-tgbot

@@ -8,6 +8,6 @@ RUN apk update \
     && apk add gcc python3-dev musl-dev libffi-dev openssl-dev make cargo
 
 COPY . .
-WORKDIR tgbot/tgbot
+WORKDIR tgbot
 RUN pip install --upgrade pip; pip install poetry; poetry config virtualenvs.create false; poetry install
-CMD poetry shell; python bot.py
+CMD poetry run tgbot

apps/tgbot/pyproject.toml

@@ -17,10 +17,12 @@ tortoise-orm = "^0.16.20"
 aiomysql = "^0.0.21"
 
 [tool.poetry.dev-dependencies]
+pytest = "^6.2.2"
+flake8 = "^3.8.4"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry.scripts]
-test = 'tgbot.test:run_all'
+tgbot = "tgbot.bot:main"

apps/tgbot/test/__init__.py

@@ -0,0 +1 @@
+

apps/tgbot/test/test_port_parsers.py

@@ -0,0 +1,74 @@
+from unittest import TestCase
+
+from tgbot.handlers.default.tcp import TCPCheckerHandler
+from tgbot.handlers.base import process_args_for_host_port,\
+                            NotEnoughArgs, InvalidPort
+
+
+class TestArgsProc(TestCase):
+
+    def test_exceptions(self):
+        """Test exceptions being raised
+        on invalid commands
+        """
+        cases = [
+            ('/cmd', NotEnoughArgs),
+            ('/cmd example.com testsarenice', InvalidPort)
+        ]
+        for cmd, exc in cases:
+            with self.subTest(command=cmd):
+                self.assertRaises(
+                    exc,
+                    process_args_for_host_port, cmd, 443
+                )
+    
+    def test_host_port(self):
+        """Test that host and port are parsed correctly
+        """
+        cases = [
+            ('/cmd example.com', 'example.com', 443),
+            ('/cmd example.com 42', 'example.com', 42),
+            ('/cmd example.com:42', 'example.com', 42)
+        ]
+
+        for cmd, host, port in cases:
+            with self.subTest(cmd=cmd, host=host, port=port):
+                test_host, test_port = process_args_for_host_port(cmd, 443)
+                self.assertEqual(test_host, host)
+                self.assertEqual(test_port, port)
+
+
+class TestTCPCheckerHandler(TestCase):
+    def setUp(self) -> None:
+        self.method = TCPCheckerHandler().process_args
+        return super().setUp()
+
+    def test_exceptions(self):
+        """Test all appropriate excpetions are raised.
+        """
+        cases = [
+            ('/cmd', NotEnoughArgs),
+            ('/cmd example.com', NotEnoughArgs),
+            ('/cmd example.com jdbnjsbndjsd', InvalidPort)
+        ]
+
+        for cmd, exc in cases:
+            with self.subTest(cmd=cmd):
+                self.assertRaises(
+                    exc,
+                    self.method, cmd
+                )
+
+    def test_host_port(self):
+        """Test that host and port are parsed correctly
+        """
+        cases = [
+            ('/cmd example.com 42', 'example.com', 42),
+            ('/cmd example.com:65', 'example.com', 65)
+        ]
+
+        for cmd, host, port in cases:
+            with self.subTest(cmd=cmd, host=host, port=port):
+                test_host, test_port = self.method(cmd)
+                self.assertEqual(test_host, host)
+                self.assertEqual(test_port, port)

apps/tgbot/tgbot/bot.py

@@ -1,12 +1,14 @@
+from asyncio import sleep
+
 from aiogram import Bot, Dispatcher, executor
 from aiogram.contrib.fsm_storage.memory import MemoryStorage
-from tgbot.middlewares import UserMiddleware, WriteCommandMetric, LoggingMiddleware, ThrottlingMiddleware
+from loguru import logger
 from tortoise import Tortoise
 from tortoise.exceptions import DBConnectionError
-from loguru import logger
-from asyncio import sleep
-import config
-import handlers
+
+from . import config, handlers
+from .middlewares import (LoggingMiddleware, ThrottlingMiddleware,
+                          UserMiddleware, WriteCommandMetric)
 
 storage = MemoryStorage()
 telegram_bot = Bot(token=config.TELEGRAM_BOT_TOKEN)
@@ -42,5 +44,10 @@ async def on_startup(disp: Dispatcher):
     disp.middleware.setup(LoggingMiddleware())
     disp.middleware.setup(UserMiddleware())
 
-if __name__ == '__main__':
+
+def main():
     executor.start_polling(dp, skip_updates=True, on_startup=on_startup)
+
+
+if __name__ == '__main__':
+    main()

apps/tgbot/tgbot/handlers/base.py

@@ -1,18 +1,19 @@
-from aiogram.types import Message
-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
-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
+from typing import Any, List, Tuple
+from uuid import uuid4
+
+from aiogram.bot import Bot
+from aiogram.types import Message
+from core.coretypes import APINodeInfo
+from httpx import Response
+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 .validators import BaseValidator, LocalhostValidator
 
 header = "Отчет о проверке хоста:" \
          "\n\n— Хост: {target_fq}"\
@@ -28,10 +29,10 @@ class SimpleCommandHandler:
     async def handler(self, message: Message):
         pass
 
-    async def process_args(self, text: str) -> list:
+    def process_args(self, text: str) -> list:
         raise NotImplemented
 
-    async def validate_target(self, target: str):
+    def validate_target(self, target: str):
         for validator in self.validators:
             validator.validate(target)
 
@@ -81,7 +82,7 @@ class CheckerTargetPortHandler(CheckerBaseHandler):
     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)
+            args = 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")
@@ -89,7 +90,7 @@ class CheckerTargetPortHandler(CheckerBaseHandler):
             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])
+            self.validate_target(args[0])
         except ValueError:  # For ip check
             pass
         except LocalhostForbidden:
@@ -102,14 +103,32 @@ class CheckerTargetPortHandler(CheckerBaseHandler):
         )
 
 
-def process_args_for_host_port(text: str, default_port: int) -> list:
+def parse_host_port(text: str, default_port: int) -> Tuple[str, int]:
+    """Parse host:port
+    """
+    text = text.strip()
     port = default_port
+    host = text
+    if ":" in text:
+        host, port = text.rsplit(":", 1)
+    elif " " in text:
+        host, port = text.rsplit(" ", 1)
+
+    try:
+        port = int(port)
+        # !Important: Don't check range if port == default_port!
+        assert port == default_port or port in range(1, 65_536)
+    except (ValueError, AssertionError):
+        raise InvalidPort(port)
+    
+    return (host, port)
+
+
+def process_args_for_host_port(text: str, default_port: int) -> Tuple[str, int]:
+    """Parse target from command
+    """
     args = text.split(' ', 1)
     if len(args) != 2:
         raise NotEnoughArgs()
-    host = args[1]
-    if ":" in host:
-        host, port = host.rsplit(":", 1)
-    elif " " in host:
-        host, port = host.rsplit(" ", 1)
-    return [host, port]
+    target = args[1]
+    return parse_host_port(target, default_port)

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

@@ -1,12 +1,12 @@
 from aiogram import Dispatcher
 
+from .icmp import ICMPCheckerHandler
+from .ipcalc import IPCalcCommandHandler
+from .minecraft import MinecraftCheckerHandler
 from .start import start_cmd
+from .tcp import TCPCheckerHandler
 from .web import WebCheckerHandler
 from .whois import WhoisCommandHandler
-from .icmp import ICMPCheckerHandler
-from .tcp import TCPCheckerHandler
-from .minecraft import MinecraftCheckerHandler
-from .ipcalc import IPCalcCommandHandler
 
 
 def setup(dp: Dispatcher):

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

@@ -1,9 +1,10 @@
 from aiogram.types import Message
-from httpx import Response
 from core.coretypes import ErrorPayload, ICMPCheckerResponse, ResponseStatus
-from ..base import CheckerBaseHandler, NotEnoughArgs, LocalhostForbidden
+from httpx import Response
+
+from ...middlewares.throttling import rate_limit
+from ..base import CheckerBaseHandler, LocalhostForbidden, NotEnoughArgs
 from ..metrics import push_status_metric
-from tgbot.middlewares.throttling import rate_limit
 
 icmp_help_message = """
 ❓ Производит проверку хоста по протоколу ICMP.
@@ -23,7 +24,7 @@ class ICMPCheckerHandler(CheckerBaseHandler):
     @rate_limit
     async def handler(self, message: Message):
         try:
-            args = await self.process_args(message.text)
+            args = self.process_args(message.text)
         except NotEnoughArgs:
             return await message.answer(icmp_help_message, parse_mode="Markdown")
         except LocalhostForbidden:
@@ -31,13 +32,13 @@ class ICMPCheckerHandler(CheckerBaseHandler):
         else:
             await self.check(message.chat.id, message.bot, dict(target=args[0], target_fq=args[0]))
 
-    async def process_args(self, text: str) -> list:
+    def process_args(self, text: str) -> list:
         args = text.split()
         if len(args) == 1:
             raise NotEnoughArgs()
         if len(args) >= 2:
             target = args[1]
-            await self.validate_target(target)
+            self.validate_target(target)
             return [target]
 
     async def prepare_message(self, res: Response):

apps/tgbot/tgbot/handlers/default/ipcalc.py

@@ -1,9 +1,10 @@
-from aiogram.types import Message
-from typing import Union
 import ipaddress
+from typing import Union
 
-from tgbot.handlers.base import SimpleCommandHandler, NotEnoughArgs
-from tgbot.middlewares.throttling import rate_limit
+from aiogram.types import Message
+
+from ...middlewares.throttling import rate_limit
+from ..base import NotEnoughArgs, SimpleCommandHandler
 
 ipcalc_help_message = """
 ❓ Калькулятор IP подсетей.
@@ -24,7 +25,7 @@ class IPCalcCommandHandler(SimpleCommandHandler):
     @rate_limit
     async def handler(self, message: Message):
         try:
-            args = await self.process_args(message.text)
+            args = self.process_args(message.text)
             network = ipaddress.ip_network(args[1], False)
         except NotEnoughArgs:
             await message.answer(self.help_message, parse_mode='Markdown')
@@ -34,7 +35,7 @@ class IPCalcCommandHandler(SimpleCommandHandler):
             msg = await self.prepare_message(network)
             await message.answer(msg)
 
-    async def process_args(self, text: str) -> list:
+    def process_args(self, text: str) -> list:
         args = text.split()
         if len(args) == 1:
             raise NotEnoughArgs

apps/tgbot/tgbot/handlers/default/minecraft.py

@@ -1,8 +1,8 @@
-from core.coretypes import ResponseStatus, ErrorPayload, MinecraftResponse
+from core.coretypes import ErrorPayload, MinecraftResponse, ResponseStatus
 from httpx import Response
 
-from tgbot.handlers.base import CheckerTargetPortHandler, process_args_for_host_port
-from tgbot.handlers.metrics import push_status_metric
+from ..base import CheckerTargetPortHandler, process_args_for_host_port
+from ..metrics import push_status_metric
 
 minecraft_help_message = """
 ❓ Получает статистику о Minecraft сервере
@@ -24,7 +24,7 @@ class MinecraftCheckerHandler(CheckerTargetPortHandler):
     def __init__(self):
         super().__init__()
 
-    async def process_args(self, text: str) -> list:
+    def process_args(self, text: str) -> list:
         return process_args_for_host_port(text, 25565)
 
     async def prepare_message(self, res: Response):

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

@@ -1,7 +1,8 @@
 from aiogram.types import Message
-from tgbot.models.user import User
-from tgbot.middlewares.throttling import rate_limit
-from tgbot.middlewares.userdata import userdata_required
+
+from ...middlewares.throttling import rate_limit
+from ...middlewares.userdata import userdata_required
+from ...models.user import User
 
 start_message = """
 

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

@@ -1,11 +1,13 @@
+from typing import Tuple
+
 from aiogram.types import Message
-from core.coretypes import ResponseStatus, ErrorPayload, PortResponse
+from core.coretypes import ErrorPayload, PortResponse, ResponseStatus
 from httpx import Response
 
-from tgbot.handlers.base import CheckerTargetPortHandler, NotEnoughArgs, InvalidPort
-from tgbot.handlers.helpers import check_int
-from tgbot.handlers.metrics import push_status_metric
-from tgbot.middlewares.throttling import rate_limit
+from ...middlewares.throttling import rate_limit
+from ..base import (CheckerTargetPortHandler, InvalidPort, NotEnoughArgs,
+                    parse_host_port)
+from ..metrics import push_status_metric
 
 tcp_help_message = """
 ❓ Производит проверку TCP порта, открыт ли он или нет
@@ -29,20 +31,15 @@ class TCPCheckerHandler(CheckerTargetPortHandler):
     async def handler(self, message: Message):
         await super(TCPCheckerHandler, self).handler(message)
 
-    async def process_args(self, text: str) -> list:
+    def process_args(self, text: str) -> Tuple[str, int]:
         args = text.split(' ', 1)
         if len(args) != 2:
             raise NotEnoughArgs()
         host = args[1]
-        if ":" in host:
-            host, port = host.rsplit(":", 1)
-        elif " " in host:
-            host, port = host.split(maxsplit=1)
-        else:
+        host, port = parse_host_port(host, -1)
+        if port == -1:
             raise NotEnoughArgs()
-        if not check_int(port):
-            raise InvalidPort()
-        return [host, port]
+        return (host, port)
 
     async def prepare_message(self, res: Response):
         message, status = await self.message_std_vals(res)

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

@@ -1,5 +1,7 @@
+from core.coretypes import (HTTP_EMOJI, ErrorPayload, HttpCheckerResponse,
+                            ResponseStatus)
 from httpx import Response
-from core.coretypes import ResponseStatus, HTTP_EMOJI, HttpCheckerResponse, ErrorPayload
+
 from ..base import CheckerTargetPortHandler, process_args_for_host_port
 from ..metrics import push_status_metric
 
@@ -22,7 +24,7 @@ class WebCheckerHandler(CheckerTargetPortHandler):
     def __init__(self):
         super().__init__()
 
-    async def process_args(self, text: str) -> list:
+    def process_args(self, text: str) -> list:
         return process_args_for_host_port(text, 80)
 
     async def prepare_message(self, res: Response):

apps/tgbot/tgbot/handlers/default/whois.py

@@ -1,15 +1,16 @@
+from dataclasses import dataclass
 from typing import Optional
 
-from whois import whois, parser
 from aiogram.types import Message
-from dataclasses import dataclass
 from whois_vu.api import WhoisSource
 from whois_vu.errors import IncorrectZone, QueryNotMatchRegexp
 
-from tgbot.handlers.whois_zones import ZONES
-from tgbot.handlers.base import SimpleCommandHandler
-from tgbot.handlers.errors import NotEnoughArgs, LocalhostForbidden
-from tgbot.middlewares.throttling import rate_limit
+from whois import parser, whois
+
+from ...middlewares.throttling import rate_limit
+from ..base import SimpleCommandHandler
+from ..errors import LocalhostForbidden, NotEnoughArgs
+from ..whois_zones import ZONES
 
 whois_help_message = """
 ❓ Вернёт информацию о домене.
@@ -116,7 +117,7 @@ class WhoisCommandHandler(SimpleCommandHandler):
     @rate_limit
     async def handler(self, message: Message):
         try:
-            args = await self.process_args(message.text)
+            args = self.process_args(message.text)
         except NotEnoughArgs:
             await message.answer(no_domain_text, parse_mode='Markdown')
         except LocalhostForbidden:
@@ -124,13 +125,13 @@ class WhoisCommandHandler(SimpleCommandHandler):
         else:
             await message.answer(create_whois_message(args[0]), parse_mode='html')
 
-    async def process_args(self, text: str) -> list:
+    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)
+            self.validate_target(host)
             return [host]  # only domain name
 
     async def prepare_message(self) -> str:

apps/tgbot/tgbot/handlers/helpers.py

@@ -1,23 +1,16 @@
-from httpx import AsyncClient, Timeout, Response
-from typing import List, Callable
-from core.coretypes import APINode
-from ipaddress import ip_address
-from contextlib import suppress
-from loguru import logger
-from aiogram.bot import Bot
-from tgbot.handlers.metrics import push_api_request_status
-from tgbot.config import NOTIFICATION_BOT_TOKEN, NOTIFICATION_USERS
-from traceback import format_exc
 import asyncio
+from contextlib import suppress
+from ipaddress import ip_address
+from traceback import format_exc
+from typing import Callable, List
 
+from aiogram.bot import Bot
+from core.coretypes import APINode
+from httpx import AsyncClient, Response, Timeout
+from loguru import logger
 
-def check_int(value) -> bool:
-    try:
-        int(value)
-    except ValueError:
-        return False
-    else:
-        return True
+from ..config import NOTIFICATION_BOT_TOKEN, NOTIFICATION_USERS
+from .metrics import push_api_request_status
 
 
 async def send_api_request(client: AsyncClient, endpoint: str, data: dict, node: APINode):

apps/tgbot/tgbot/handlers/metrics.py

@@ -1,6 +1,9 @@
-from aioinflux import InfluxDBClient
 from typing import Dict
-from tgbot.config import INFLUX_DB, INFLUX_HOST, INFLUX_PORT, INFLUX_PASSWORD, INFLUX_USERNAME
+
+from aioinflux import InfluxDBClient
+
+from ..config import (INFLUX_DB, INFLUX_HOST, INFLUX_PASSWORD, INFLUX_PORT,
+                      INFLUX_USERNAME)
 
 
 async def push_metric(measurement, tags: Dict, fields: Dict):

apps/tgbot/tgbot/handlers/validators.py

@@ -1,6 +1,7 @@
-from .errors import LocalhostForbidden
-from ipaddress import ip_address
 from contextlib import suppress
+from ipaddress import ip_address
+
+from .errors import LocalhostForbidden
 
 
 class BaseValidator:

apps/tgbot/tgbot/middlewares/__init__.py

@@ -1,4 +1,4 @@
-from tgbot.middlewares.write_command_metric import WriteCommandMetric
-from tgbot.middlewares.logging import LoggingMiddleware
-from tgbot.middlewares.throttling import ThrottlingMiddleware
-from tgbot.middlewares.userdata import UserMiddleware
+from .logging import LoggingMiddleware
+from .throttling import ThrottlingMiddleware
+from .userdata import UserMiddleware
+from .write_command_metric import WriteCommandMetric

apps/tgbot/tgbot/middlewares/throttling.py

@@ -1,9 +1,10 @@
+import asyncio
+
 from aiogram import Dispatcher, types
 from aiogram.dispatcher import DEFAULT_RATE_LIMIT
 from aiogram.dispatcher.handler import CancelHandler, current_handler
 from aiogram.dispatcher.middlewares import BaseMiddleware
 from aiogram.utils.exceptions import Throttled
-import asyncio
 
 
 def rate_limit(func):

apps/tgbot/tgbot/middlewares/userdata.py

@@ -1,8 +1,8 @@
 from aiogram.dispatcher.handler import current_handler
 from aiogram.dispatcher.middlewares import BaseMiddleware
-from aiogram.types import Message, CallbackQuery
+from aiogram.types import CallbackQuery, Message
 
-from tgbot.models import User
+from ..models import User
 
 
 def userdata_required(func):

apps/tgbot/tgbot/middlewares/write_command_metric.py

@@ -1,7 +1,8 @@
 from aiogram.dispatcher.middlewares import BaseMiddleware
 from aiogram.types import Message
-from tgbot.handlers.metrics import push_metric
-from tgbot.models import User, UserCheckRequests
+
+from ..handlers.metrics import push_metric
+from ..models import User, UserCheckRequests
 
 
 class WriteCommandMetric(BaseMiddleware):

apps/tgbot/tgbot/nodes.py

@@ -1,6 +1,7 @@
-from core.coretypes import APINode
 from typing import List
 
+from core.coretypes import APINode
+
 nodes: List[APINode] = [
     APINode("http://localhost:8080", "CHANGE_TOKEN_BY_ENV"),
     APINode("http://localhost:8080", "CHANGE_TOKEN_BY_ENV"),

apps/tgbot/tgbot/test/__init__.py

@@ -1,2 +0,0 @@
-def run_all():
-    from . import port_parsers

apps/tgbot/tgbot/test/port_parsers.py

@@ -1,57 +0,0 @@
-import asyncio
-
-from ..handlers.default.tcp import TCPCheckerHandler
-from ..handlers.base import process_args_for_host_port,\
-                            NotEnoughArgs, InvalidPort
-
-
-try:
-    args = "/cmd"
-    process_args_for_host_port(args, 443)
-except NotEnoughArgs:
-    pass
-args = "/cmd example.com"
-host, port = process_args_for_host_port(args, 443)
-assert port == 443
-
-args = "/cmd example.com 42"
-host, port = process_args_for_host_port(args, 443)
-assert port == "42"  # TODO: FIX THIS SHIT
-
-args = "/cmd example.com:42"
-host, port = process_args_for_host_port(args, 443)
-assert port == "42"
-
-try:
-    args = "/cmd example.com fucktests"
-except InvalidPort:
-    pass
-
-method = TCPCheckerHandler().process_args
-
-
-async def test():
-    try:
-        args = "/cmd"
-        await method(args)
-        args = "/cmd example.com"
-        await method(args)
-    except NotEnoughArgs:
-        pass
-
-    args = "/cmd example.com 42"
-    host, port = await method(args)
-    assert port == "42"
-
-    args = "/cmd example.com:42"
-    host, port = await method(args)
-    assert port == "42"
-
-    try:
-        args = "/cmd example.com jdbnjsbndjsd"
-        await method(args)
-    except InvalidPort:
-        pass
-
-
-asyncio.run(test())

pyproject.toml

@@ -8,6 +8,8 @@ authors = ["kiriharu <kiriharu@yandex.ru>"]
 python = "^3.8.2"
 
 [tool.poetry.dev-dependencies]
+pytest = "^6.2.2"
+flake8 = "^3.8.4"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]