From a54a938c2f3ecf68005c528fe19cfc48d29c2e92 Mon Sep 17 00:00:00 2001
From: Termonoid <termonoid@yandex.ru>
Date: Sun, 31 Jan 2021 10:27:22 +0300
Subject: [PATCH] Port can be defined as host:port

---
 apps/tgbot/tgbot/handlers/base.py        | 16 +++++++---------
 apps/tgbot/tgbot/handlers/default/tcp.py | 17 ++++++++++-------
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/apps/tgbot/tgbot/handlers/base.py b/apps/tgbot/tgbot/handlers/base.py
index 623e850..c880ab0 100644
--- a/apps/tgbot/tgbot/handlers/base.py
+++ b/apps/tgbot/tgbot/handlers/base.py
@@ -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:
+    port = default_port
+    args = text.split(' ', 1)
+    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()
     host = args[1]
+    if ":" in host:
+        host, port = host.rsplit(":", 1)
+    elif " " in host:
+        host, port = host.rsplit(" ", 1)
     return [host, port]
diff --git a/apps/tgbot/tgbot/handlers/default/tcp.py b/apps/tgbot/tgbot/handlers/default/tcp.py
index a2aa1fc..14edd65 100644
--- a/apps/tgbot/tgbot/handlers/default/tcp.py
+++ b/apps/tgbot/tgbot/handlers/default/tcp.py
@@ -29,15 +29,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()
-        if len(args) >= 3:
-            port = args[2]
-            if not check_int(port):
-                raise InvalidPort()
         host = args[1]
+        if ":" in host:
+            host, port = host.rsplit(":", 1)
+        elif " " in host:
+            host, port = host.split(maxsplit=1)
+        else:
+            raise NotEnoughArgs()
+        if not check_int(port):
+            raise InvalidPort()
         return [host, port]
 
     async def prepare_message(self, res: Response):