diff options
| -rw-r--r-- | kvmd/apps/vnc/server.py | 6 | ||||
| -rw-r--r-- | kvmd/clients/kvmd.py | 91 |
2 files changed, 64 insertions, 33 deletions
diff --git a/kvmd/apps/vnc/server.py b/kvmd/apps/vnc/server.py index b860b261..12baeb20 100644 --- a/kvmd/apps/vnc/server.py +++ b/kvmd/apps/vnc/server.py @@ -233,7 +233,7 @@ class _Client(RfbClient): # pylint: disable=too-many-instance-attributes # ===== async def _authorize_userpass(self, user: str, passwd: str) -> bool: - if (await self.__kvmd.authorize(user, passwd)): + if (await self.__kvmd.auth.check(user, passwd)): self.__authorized.set_result((user, passwd)) return True return False @@ -286,7 +286,7 @@ class _Client(RfbClient): # pylint: disable=too-many-instance-attributes (user, passwd) = self.__authorized.result() get_logger(0).info("[main] Client %s: Applying streamer params: quality=%d%%; desired_fps=%d ...", self._remote, self._encodings.tight_jpeg_quality, self.__desired_fps) - await self.__kvmd.set_streamer_params(user, passwd, self._encodings.tight_jpeg_quality, self.__desired_fps) + await self.__kvmd.streamer.set_params(user, passwd, self._encodings.tight_jpeg_quality, self.__desired_fps) async def _on_fb_update_request(self) -> None: async with self.__lock: @@ -326,7 +326,7 @@ class VncServer: # pylint: disable=too-many-instance-attributes logger.info("Preparing client %s ...", remote) try: try: - none_auth_only = await kvmd.authorize("", "") + none_auth_only = await kvmd.auth.check("", "") except KvmdError as err: logger.error("Client %s: Can't check KVMD auth mode: %s", remote, err) return diff --git a/kvmd/clients/kvmd.py b/kvmd/clients/kvmd.py index b2667053..64d2ddd6 100644 --- a/kvmd/clients/kvmd.py +++ b/kvmd/clients/kvmd.py @@ -24,6 +24,7 @@ import contextlib from typing import Dict from typing import AsyncGenerator +from typing import Union import aiohttp @@ -32,12 +33,15 @@ from .. import make_user_agent # ===== class KvmdError(Exception): - def __init__(self, err: Exception): - super().__init__(f"{type(err).__name__}: {err}") + def __init__(self, err: Union[Exception, str]) -> None: + if isinstance(err, Exception): + super().__init__(f"{type(err).__name__}: {err}") + else: + super().__init__(err) # ===== -class KvmdClient: +class _BaseClientPart: def __init__( self, host: str, @@ -52,16 +56,33 @@ class KvmdClient: self.__unix_path = unix_path self.__timeout = timeout - # ===== + def _make_url(self, handle: str) -> str: + assert not handle.startswith("/"), handle + return f"http://{self.__host}:{self.__port}/{handle}" - async def authorize(self, user: str, passwd: str) -> bool: + def _make_session(self, user: str, passwd: str) -> aiohttp.ClientSession: + kwargs: Dict = { + "headers": { + "X-KVMD-User": user, + "X-KVMD-Passwd": passwd, + "User-Agent": make_user_agent("KVMD-VNC"), + }, + "timeout": aiohttp.ClientTimeout(total=self.__timeout), + } + if self.__unix_path: + kwargs["connector"] = aiohttp.UnixConnector(path=self.__unix_path) + return aiohttp.ClientSession(**kwargs) + + +class _AuthClientPart(_BaseClientPart): + async def check(self, user: str, passwd: str) -> bool: try: - async with self.__make_session(user, passwd) as session: - async with session.get(f"http://{self.__host}:{self.__port}/auth/check") as response: + async with self._make_session(user, passwd) as session: + async with session.get(self._make_url("auth/check")) as response: response.raise_for_status() if response.status == 200: return True - raise RuntimeError(f"Invalid OK response: {response.status} {await response.text()}") + raise KvmdError(f"Invalid OK response: {response.status} {await response.text()}") except aiohttp.ClientResponseError as err: if err.status in [401, 403]: return False @@ -69,37 +90,47 @@ class KvmdClient: except aiohttp.ClientError as err: raise KvmdError(err) - @contextlib.asynccontextmanager - async def ws(self, user: str, passwd: str) -> AsyncGenerator[aiohttp.ClientWebSocketResponse, None]: - try: - async with self.__make_session(user, passwd) as session: - async with session.ws_connect(f"http://{self.__host}:{self.__port}/ws") as ws: - yield ws - except aiohttp.ClientError as err: - raise KvmdError(err) - async def set_streamer_params(self, user: str, passwd: str, quality: int, desired_fps: int) -> None: +class _StreamerClientPart(_BaseClientPart): + async def set_params(self, user: str, passwd: str, quality: int, desired_fps: int) -> None: try: - async with self.__make_session(user, passwd) as session: + async with self._make_session(user, passwd) as session: async with session.post( - url=f"http://{self.__host}:{self.__port}/streamer/set_params", + url=self._make_url("streamer/set_params"), params={"quality": quality, "desired_fps": desired_fps}, ) as response: response.raise_for_status() except aiohttp.ClientError as err: raise KvmdError(err) - # ===== - def __make_session(self, user: str, passwd: str) -> aiohttp.ClientSession: +# ===== +class KvmdClient(_BaseClientPart): + def __init__( + self, + host: str, + port: int, + unix_path: str, + timeout: float, + ) -> None: + kwargs: Dict = { - "headers": { - "X-KVMD-User": user, - "X-KVMD-Passwd": passwd, - "User-Agent": make_user_agent("KVMD-VNC"), - }, - "timeout": aiohttp.ClientTimeout(total=self.__timeout), + "host": host, + "port": port, + "unix_path": unix_path, + "timeout": timeout, } - if self.__unix_path: - kwargs["connector"] = aiohttp.UnixConnector(path=self.__unix_path) - return aiohttp.ClientSession(**kwargs) + + super().__init__(**kwargs) + + self.auth = _AuthClientPart(**kwargs) + self.streamer = _StreamerClientPart(**kwargs) + + @contextlib.asynccontextmanager + async def ws(self, user: str, passwd: str) -> AsyncGenerator[aiohttp.ClientWebSocketResponse, None]: + try: + async with self._make_session(user, passwd) as session: + async with session.ws_connect(self._make_url("ws")) as ws: + yield ws + except aiohttp.ClientError as err: + raise KvmdError(err) |