Merge pull request #26 from Termonoid/main

Port can be defined as host:port

apps/tgbot/pyproject.toml

@@ -21,3 +21,6 @@ aiomysql = "^0.0.21"
 [build-system]
 requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.scripts]
+test = 'tgbot.test:run_all'

apps/tgbot/tgbot/handlers/base.py

@@ -6,7 +6,7 @@ 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
+from .helpers import send_api_requests
 from .errors import NotEnoughArgs, InvalidPort, LocalhostForbidden
 from .validators import BaseValidator, LocalhostValidator
 from tgbot.middlewares.throttling import rate_limit
@@ -103,15 +103,13 @@ class CheckerTargetPortHandler(CheckerBaseHandler):
 
 
 def process_args_for_host_port(text: str, default_port: int) -> list:
-    port = None
-    args = text.split()
-    if len(args) < 2:
-        raise NotEnoughArgs()
-    if len(args) == 2:
     port = default_port
-    if len(args) == 3:
-        port = args[2]
-        if not check_int(port):
-            raise InvalidPort()
+    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]

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

@@ -9,6 +9,7 @@ minecraft_help_message = """
 
 Использование:
  `/minecraft <hostname> <port>`
+ `/minecraft <hostname>:<port>`
  `/minecraft <hostname>` - автоматически выставит порт 25565
 """
 

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

@@ -12,6 +12,7 @@ tcp_help_message = """
 
 Использование:
  `/tcp <hostname> <port>`
+ `/tcp <hostname>:<port>`
 """
 
 invalid_port = """❗Неправильный порт. Напишите /tcp чтобы увидеть справку к данному способу проверки."""
@@ -29,15 +30,18 @@ class TCPCheckerHandler(CheckerTargetPortHandler):
         await super(TCPCheckerHandler, self).handler(message)
 
     async def process_args(self, text: str) -> list:
-        port = None
-        args = text.split()
-        if len(args) < 3:
+        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:
             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):

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

@@ -8,6 +8,7 @@ web_help_message = """
 
 Использование:
  `/web <hostname> <port>`
+ `/web <hostname>:<port>`
  `/web <hostname>` - автоматически выставит 80 порт
 """
 

apps/tgbot/tgbot/test/__init__.py

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

apps/tgbot/tgbot/test/port_parsers.py

@@ -0,0 +1,57 @@
+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())