refactoring

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)