refactoring

2 files changed, 67 insertions, 53 deletions

diff --git a/kvmd/apps/janus/runner.py b/kvmd/apps/janus/runner.py
index e08fade0..cb46562a 100644
--- a/kvmd/apps/janus/runner.py
+++ b/kvmd/apps/janus/runner.py
@@ -11,15 +11,16 @@ from ... import aioproc
 
 from ...logging import get_logger
 
+from .stun import StunNatType
 from .stun import Stun
 
 
 # =====
 @dataclasses.dataclass(frozen=True)
 class _Netcfg:
-    nat_type: str = dataclasses.field(default="")
-    src_ip: str = dataclasses.field(default="")
-    ext_ip: str = dataclasses.field(default="")
+    nat_type:  StunNatType = dataclasses.field(default=StunNatType.ERROR)
+    src_ip:    str = dataclasses.field(default="")
+    ext_ip:    str = dataclasses.field(default="")
     stun_host: str = dataclasses.field(default="")
     stun_port: int = dataclasses.field(default=0)
 
@@ -92,8 +93,9 @@ class JanusRunner:  # pylint: disable=too-many-instance-attributes
 
     async def __get_netcfg(self) -> _Netcfg:
         src_ip = (self.__get_default_ip() or "0.0.0.0")
-        (stun, (nat_type, ext_ip)) = await self.__get_stun_info(src_ip)
-        return _Netcfg(nat_type, src_ip, ext_ip, stun.host, stun.port)
+        info = await self.__stun.get_info(src_ip, 0)
+        # В текущей реализации _Netcfg() это копия StunInfo()
+        return _Netcfg(**dataclasses.asdict(info))
 
     def __get_default_ip(self) -> str:
         try:
@@ -115,13 +117,6 @@ class JanusRunner:  # pylint: disable=too-many-instance-attributes
             get_logger().error("Can't get default IP: %s", tools.efmt(err))
         return ""
 
-    async def __get_stun_info(self, src_ip: str) -> tuple[Stun, tuple[str, str]]:
-        try:
-            return (self.__stun, (await self.__stun.get_info(src_ip, 0)))
-        except Exception as err:
-            get_logger().error("Can't get STUN info: %s", tools.efmt(err))
-            return (self.__stun, ("", ""))
-
     # =====
 
     @aiotools.atomic_fg
diff --git a/kvmd/apps/janus/stun.py b/kvmd/apps/janus/stun.py
index 5fea9da6..7c37dc77 100644
--- a/kvmd/apps/janus/stun.py
+++ b/kvmd/apps/janus/stun.py
@@ -4,6 +4,7 @@ import ipaddress
 import struct
 import secrets
 import dataclasses
+import enum
 
 from ... import tools
 from ... import aiotools
@@ -12,29 +13,39 @@ from ...logging import get_logger
 
 
 # =====
+class StunNatType(enum.Enum):
+    ERROR               = ""
+    BLOCKED             = "Blocked"
+    OPEN_INTERNET       = "Open Internet"
+    SYMMETRIC_UDP_FW    = "Symmetric UDP Firewall"
+    FULL_CONE_NAT       = "Full Cone NAT"
+    RESTRICTED_NAT      = "Restricted NAT"
+    RESTRICTED_PORT_NAT = "Restricted Port NAT"
+    SYMMETRIC_NAT       = "Symmetric NAT"
+    CHANGED_ADDR_ERROR  = "Error when testing on Changed-IP and Port"
+
+
 @dataclasses.dataclass(frozen=True)
-class StunAddress:
-    ip: str
+class StunInfo:
+    nat_type:  StunNatType
+    src_ip:    str
+    ext_ip:    str
+    stun_host: str
+    stun_port: int
+
+
[email protected](frozen=True)
+class _StunAddress:
+    ip:   str
     port: int
 
 
 @dataclasses.dataclass(frozen=True)
-class StunResponse:
-    ok: bool
-    ext: (StunAddress | None) = dataclasses.field(default=None)
-    src: (StunAddress | None) = dataclasses.field(default=None)
-    changed: (StunAddress | None) = dataclasses.field(default=None)
-
-
-class StunNatType:
-    BLOCKED = "Blocked"
-    OPEN_INTERNET = "Open Internet"
-    SYMMETRIC_UDP_FW = "Symmetric UDP Firewall"
-    FULL_CONE_NAT = "Full Cone NAT"
-    RESTRICTED_NAT = "Restricted NAT"
-    RESTRICTED_PORT_NAT = "Restricted Port NAT"
-    SYMMETRIC_NAT = "Symmetric NAT"
-    CHANGED_ADDR_ERROR = "Error when testing on Changed-IP and Port"
+class _StunResponse:
+    ok:      bool
+    ext:     (_StunAddress | None) = dataclasses.field(default=None)
+    src:     (_StunAddress | None) = dataclasses.field(default=None)
+    changed: (_StunAddress | None) = dataclasses.field(default=None)
 
 
 # =====
@@ -50,33 +61,44 @@ class Stun:
         retries_delay: float,
     ) -> None:
 
-        self.host = host
-        self.port = port
+        self.__host = host
+        self.__port = port
         self.__timeout = timeout
         self.__retries = retries
         self.__retries_delay = retries_delay
 
         self.__sock: (socket.socket | None) = None
 
-    async def get_info(self, src_ip: str, src_port: int) -> tuple[str, str]:
+    async def get_info(self, src_ip: str, src_port: int) -> StunInfo:
         (family, _, _, _, addr) = socket.getaddrinfo(src_ip, src_port, type=socket.SOCK_DGRAM)[0]
+        nat_type = StunNatType.ERROR
+        ext_ip = ""
         try:
             with socket.socket(family, socket.SOCK_DGRAM) as self.__sock:
                 self.__sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                 self.__sock.settimeout(self.__timeout)
                 self.__sock.bind(addr)
                 (nat_type, response) = await self.__get_nat_type(src_ip)
-                return (nat_type, (response.ext.ip if response.ext is not None else ""))
+                ext_ip = (response.ext.ip if response.ext is not None else "")
+        except Exception as err:
+            get_logger(0).error("Can't get STUN info: %s", tools.efmt(err))
         finally:
             self.__sock = None
-
-    async def __get_nat_type(self, src_ip: str) -> tuple[str, StunResponse]:  # pylint: disable=too-many-return-statements
-        first = await self.__make_request("First probe")
+        return StunInfo(
+            nat_type=nat_type,
+            src_ip=src_ip,
+            ext_ip=ext_ip,
+            stun_host=self.__host,
+            stun_port=self.__port,
+        )
+
+    async def __get_nat_type(self, src_ip: str) -> tuple[StunNatType, _StunResponse]:  # pylint: disable=too-many-return-statements
+        first = await self.__make_request("First probe", self.__host, b"")
         if not first.ok:
             return (StunNatType.BLOCKED, first)
 
         request = struct.pack(">HHI", 0x0003, 0x0004, 0x00000006)  # Change-Request
-        response = await self.__make_request("Change request [ext_ip == src_ip]", request)
+        response = await self.__make_request("Change request [ext_ip == src_ip]", self.__host, request)
 
         if first.ext is not None and first.ext.ip == src_ip:
             if response.ok:
@@ -88,20 +110,20 @@ class Stun:
 
         if first.changed is None:
             raise RuntimeError(f"Changed addr is None: {first}")
-        response = await self.__make_request("Change request [ext_ip != src_ip]", addr=first.changed)
+        response = await self.__make_request("Change request [ext_ip != src_ip]", first.changed, b"")
         if not response.ok:
             return (StunNatType.CHANGED_ADDR_ERROR, response)
 
         if response.ext == first.ext:
             request = struct.pack(">HHI", 0x0003, 0x0004, 0x00000002)
-            response = await self.__make_request("Change port", request, addr=first.changed.ip)
+            response = await self.__make_request("Change port", first.changed.ip, request)
             if response.ok:
                 return (StunNatType.RESTRICTED_NAT, response)
             return (StunNatType.RESTRICTED_PORT_NAT, response)
 
         return (StunNatType.SYMMETRIC_NAT, response)
 
-    async def __make_request(self, ctx: str, request: bytes=b"", addr: (StunAddress | str | None)=None) -> StunResponse:
+    async def __make_request(self, ctx: str, addr: (_StunAddress | str), request: bytes) -> _StunResponse:
         # TODO: Support IPv6 and RFC 5389
         # The first 4 bytes of the response are the Type (2) and Length (2)
         # The 5th byte is Reserved
@@ -111,13 +133,10 @@ class Stun:
         # More info at: https://tools.ietf.org/html/rfc3489#section-11.2.1
         # And at: https://tools.ietf.org/html/rfc5389#section-15.1
 
-        if isinstance(addr, StunAddress):
+        if isinstance(addr, _StunAddress):
             addr_t = (addr.ip, addr.port)
-        elif isinstance(addr, str):
-            addr_t = (addr, self.port)
-        else:
-            assert addr is None
-            addr_t = (self.host, self.port)
+        else:  # str
+            addr_t = (addr, self.__port)
 
         # https://datatracker.ietf.org/doc/html/rfc5389#section-6
         trans_id = b"\x21\x12\xA4\x42" + secrets.token_bytes(12)
@@ -130,9 +149,9 @@ class Stun:
         if error:
             get_logger(0).error("%s: Can't perform STUN request after %d retries; last error: %s",
                                 ctx, self.__retries, error)
-            return StunResponse(ok=False)
+            return _StunResponse(ok=False)
 
-        parsed: dict[str, StunAddress] = {}
+        parsed: dict[str, _StunAddress] = {}
         offset = 0
         remaining = len(response)
         while remaining > 0:
@@ -148,7 +167,7 @@ class Stun:
                 parsed[field] = self.__parse_address(response[offset:], (trans_id if attr_type == 0x0020 else b""))
             offset += attr_len
             remaining -= (4 + attr_len)
-        return StunResponse(ok=True, **parsed)
+        return _StunResponse(ok=True, **parsed)
 
     async def __inner_make_request(self, trans_id: bytes, request: bytes, addr: tuple[str, int]) -> tuple[bytes, str]:
         assert self.__sock is not None
@@ -172,13 +191,13 @@ class Stun:
 
         return (response[20 : 20 + payload_len], "")  # noqa: E203
 
-    def __parse_address(self, data: bytes, trans_id: bytes) -> StunAddress:
+    def __parse_address(self, data: bytes, trans_id: bytes) -> _StunAddress:
         family = data[1]
         port = struct.unpack(">H", self.__trans_xor(data[2:4], trans_id))[0]
         if family == 0x01:
-            return StunAddress(str(ipaddress.IPv4Address(self.__trans_xor(data[4:8], trans_id))), port)
+            return _StunAddress(str(ipaddress.IPv4Address(self.__trans_xor(data[4:8], trans_id))), port)
         elif family == 0x02:
-            return StunAddress(str(ipaddress.IPv6Address(self.__trans_xor(data[4:20], trans_id))), port)
+            return _StunAddress(str(ipaddress.IPv6Address(self.__trans_xor(data[4:20], trans_id))), port)
         raise RuntimeError(f"Unknown family; received: {family}")
 
     def __trans_xor(self, data: bytes, trans_id: bytes) -> bytes: