GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2025-12-12 01:00:29 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactoring

Browse Source
This commit is contained in:
Maxim Devaev 2024-09-16 23:07:38 +03:00
parent b779c18530
commit c57334f214
3 changed files with 68 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
kvmd/apps/janus/runner.py
View File

@ -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

99
kvmd/apps/janus/stun.py
View File

@ -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
@dataclasses.dataclass(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
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[str, StunResponse]: # pylint: disable=too-many-return-statements
first = await self.__make_request("First probe")
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:

3
testenv/linters/flake8.ini
View File

@ -1,8 +1,9 @@
[flake8]
inline-quotes = double
max-line-length = 160
ignore = W503, E227, E241, E252, Q003
ignore = W503, E221, E227, E241, E252, Q003
# W503 line break before binary operator
# E221 multiple spaces before operator
# E227 missing whitespace around bitwise or shift operator
# E241 multiple spaces after
# E252 missing whitespace around parameter equals