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

vnc: additional auth check

Browse Source
This commit is contained in:
Maxim Devaev 2025-05-07 12:46:45 +03:00
parent 8d7f89e8f1
commit bb0656c0cb
1 changed files with 15 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
kvmd/apps/vnc/rfb/__init__.py
View File

@ -46,6 +46,11 @@ from .crypto import rfb_encrypt_challenge
from .stream import RfbClientStream from .stream import RfbClientStream
# =====
class _SecurityError(Exception):
pass
# ===== # =====
class RfbClient(RfbClientStream): # pylint: disable=too-many-instance-attributes class RfbClient(RfbClientStream): # pylint: disable=too-many-instance-attributes
# https://github.com/rfbproto/rfbproto/blob/master/rfbproto.rst # https://github.com/rfbproto/rfbproto/blob/master/rfbproto.rst
@ -94,6 +99,8 @@ class RfbClient(RfbClientStream): # pylint: disable=too-many-instance-attribute
self.__fb_cont_updates = False self.__fb_cont_updates = False
self.__fb_reset_h264 = False self.__fb_reset_h264 = False
self.__authorized = False
self.__lock = asyncio.Lock() self.__lock = asyncio.Lock()
# ===== # =====
@ -134,6 +141,8 @@ class RfbClient(RfbClientStream): # pylint: disable=too-many-instance-attribute
async def __main_task_loop(self) -> None: async def __main_task_loop(self) -> None:
await self.__handshake_version() await self.__handshake_version()
await self.__handshake_security() await self.__handshake_security()
if not self.__authorized:
raise _SecurityError()
await self.__handshake_init() await self.__handshake_init()
await self.__main_loop() await self.__main_loop()
@ -385,6 +394,7 @@ class RfbClient(RfbClientStream): # pylint: disable=too-many-instance-attribute
if allow: if allow:
get_logger(0).info("%s [main]: %s", self._remote, allow_msg) get_logger(0).info("%s [main]: %s", self._remote, allow_msg)
await self._write_struct("access OK", "L", 0) await self._write_struct("access OK", "L", 0)
self.__authorized = True
else: else:
await self._write_struct("access denial flag", "L", 1, drain=(self.__rfb_version < 8)) await self._write_struct("access denial flag", "L", 1, drain=(self.__rfb_version < 8))
if self.__rfb_version >= 8: if self.__rfb_version >= 8:
@ -394,6 +404,9 @@ class RfbClient(RfbClientStream): # pylint: disable=too-many-instance-attribute
# ===== # =====
async def __handshake_init(self) -> None: async def __handshake_init(self) -> None:
if not self.__authorized:
raise _SecurityError()
await self._read_number("initial shared flag", "B") # Shared flag, ignored await self._read_number("initial shared flag", "B") # Shared flag, ignored
await self._write_struct("initial FB size", "HH", self._width, self._height, drain=False) await self._write_struct("initial FB size", "HH", self._width, self._height, drain=False)
@ -417,6 +430,8 @@ class RfbClient(RfbClientStream): # pylint: disable=too-many-instance-attribute
# ===== # =====
async def __main_loop(self) -> None: async def __main_loop(self) -> None:
if not self.__authorized:
raise _SecurityError()
handlers = { handlers = {
0: self.__handle_set_pixel_format, 0: self.__handle_set_pixel_format,
2: self.__handle_set_encodings, 2: self.__handle_set_encodings,