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

refactoring

Browse Source
This commit is contained in:
Maxim Devaev 2022-11-20 19:52:33 +03:00
parent 299fadf0fa
commit d1800561e6
1 changed files with 14 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -88,20 +88,20 @@ class Stun:
if first.changed is None: if first.changed is None:
raise RuntimeError(f"Changed addr is None: {first}") raise RuntimeError(f"Changed addr is None: {first}")
response = await self.__make_request("Change request [ext_ip != src_ip]", b"", first.changed.ip, first.changed.port) response = await self.__make_request("Change request [ext_ip != src_ip]", addr=first.changed)
if not response.ok: if not response.ok:
return (StunNatType.CHANGED_ADDR_ERROR, response) return (StunNatType.CHANGED_ADDR_ERROR, response)
if response.ext == first.ext: if response.ext == first.ext:
request = struct.pack(">HHI", 0x0003, 0x0004, 0x00000002) request = struct.pack(">HHI", 0x0003, 0x0004, 0x00000002)
response = await self.__make_request("Change port", request, first.changed.ip) response = await self.__make_request("Change port", request, addr=first.changed.ip)
if response.ok: if response.ok:
return (StunNatType.RESTRICTED_NAT, response) return (StunNatType.RESTRICTED_NAT, response)
return (StunNatType.RESTRICTED_PORT_NAT, response) return (StunNatType.RESTRICTED_PORT_NAT, response)
return (StunNatType.SYMMETRIC_NAT, response) return (StunNatType.SYMMETRIC_NAT, response)
async def __make_request(self, ctx: str, request: bytes=b"", host: str="", port: int=0) -> StunResponse: async def __make_request(self, ctx: str, request: bytes=b"", addr: (StunAddress | str | None)=None) -> StunResponse:
# TODO: Support IPv6 and RFC 5389 # TODO: Support IPv6 and RFC 5389
# The first 4 bytes of the response are the Type (2) and Length (2) # The first 4 bytes of the response are the Type (2) and Length (2)
# The 5th byte is Reserved # The 5th byte is Reserved
@ -111,11 +111,19 @@ class Stun:
# More info at: https://tools.ietf.org/html/rfc3489#section-11.2.1 # More info at: https://tools.ietf.org/html/rfc3489#section-11.2.1
# And at: https://tools.ietf.org/html/rfc5389#section-15.1 # And at: https://tools.ietf.org/html/rfc5389#section-15.1
if isinstance(addr, StunAddress):
addr_t = (addr.ip, addr.port)
elif isinstance(addr, str):
addr_t = (addr, self.port)
else:
assert addr_t is None
addr_t = (self.host, self.port)
# https://datatracker.ietf.org/doc/html/rfc5389#section-6 # https://datatracker.ietf.org/doc/html/rfc5389#section-6
trans_id = b"\x21\x12\xA4\x42" + secrets.token_bytes(12) trans_id = b"\x21\x12\xA4\x42" + secrets.token_bytes(12)
(response, error) = (b"", "") (response, error) = (b"", "")
for _ in range(self.__retries): for _ in range(self.__retries):
(response, error) = await self.__inner_make_request(trans_id, request, host, port) (response, error) = await self.__inner_make_request(trans_id, request, addr_t)
if not error: if not error:
break break
await asyncio.sleep(self.__retries_delay) await asyncio.sleep(self.__retries_delay)
@ -142,13 +150,13 @@ class Stun:
remaining -= (4 + 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, host: str, port: int) -> tuple[bytes, str]: async def __inner_make_request(self, trans_id: bytes, request: bytes, addr: tuple[str, int]) -> tuple[bytes, str]:
assert self.__sock is not None assert self.__sock is not None
request = struct.pack(">HH", 0x0001, len(request)) + trans_id + request # Bind Request request = struct.pack(">HH", 0x0001, len(request)) + trans_id + request # Bind Request
try: try:
await aiotools.run_async(self.__sock.sendto, request, ((host or self.host), (port or self.port))) await aiotools.run_async(self.__sock.sendto, request, addr)
except Exception as err: except Exception as err:
return (b"", f"Send error: {tools.efmt(err)}") return (b"", f"Send error: {tools.efmt(err)}")
try: try: