refactoring

1 files changed, 41 insertions, 34 deletions

diff --git a/kvmd/clients/kvmd.py b/kvmd/clients/kvmd.py
index 64d2ddd6..0b33eb42 100644
--- a/kvmd/clients/kvmd.py
+++ b/kvmd/clients/kvmd.py
@@ -24,20 +24,10 @@ import contextlib
 
 from typing import Dict
 from typing import AsyncGenerator
-from typing import Union
 
 import aiohttp
 
-from .. import make_user_agent
-
-
-# =====
-class KvmdError(Exception):
-    def __init__(self, err: Union[Exception, str]) -> None:
-        if isinstance(err, Exception):
-            super().__init__(f"{type(err).__name__}: {err}")
-        else:
-            super().__init__(err)
+from .. import aiotools
 
 
 # =====
@@ -48,6 +38,7 @@ class _BaseClientPart:
         port: int,
         unix_path: str,
         timeout: float,
+        user_agent: str,
     ) -> None:
 
         assert port or unix_path
@@ -55,17 +46,14 @@ class _BaseClientPart:
         self.__port = port
         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}"
+        self.__user_agent = user_agent
 
     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"),
+                "User-Agent": self.__user_agent,
             },
             "timeout": aiohttp.ClientTimeout(total=self.__timeout),
         }
@@ -73,35 +61,54 @@ class _BaseClientPart:
             kwargs["connector"] = aiohttp.UnixConnector(path=self.__unix_path)
         return aiohttp.ClientSession(**kwargs)
 
+    def _make_url(self, handle: str) -> str:
+        assert not handle.startswith("/"), handle
+        return f"http://{self.__host}:{self.__port}/{handle}"
+
 
 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(self._make_url("auth/check")) as response:
-                    response.raise_for_status()
-                    if response.status == 200:
-                        return True
-                    raise KvmdError(f"Invalid OK response: {response.status} {await response.text()}")
+                    aiotools.raise_not_200(response)
+                    return True
         except aiohttp.ClientResponseError as err:
             if err.status in [401, 403]:
                 return False
-            raise KvmdError(err)
-        except aiohttp.ClientError as err:
-            raise KvmdError(err)
+            raise
 
 
 class _StreamerClientPart(_BaseClientPart):
     async def set_params(self, user: str, passwd: str, quality: int, desired_fps: int) -> None:
+        async with self._make_session(user, passwd) as session:
+            async with session.post(
+                url=self._make_url("streamer/set_params"),
+                params={"quality": quality, "desired_fps": desired_fps},
+            ) as response:
+                aiotools.raise_not_200(response)
+
+
+class _AtxClientPart(_BaseClientPart):
+    async def get_state(self, user: str, passwd: str) -> Dict:
+        async with self._make_session(user, passwd) as session:
+            async with session.get(self._make_url("atx")) as response:
+                aiotools.raise_not_200(response)
+                return (await response.json())["result"]
+
+    async def switch_power(self, user: str, passwd: str, action: str) -> bool:
         try:
             async with self._make_session(user, passwd) as session:
                 async with session.post(
-                    url=self._make_url("streamer/set_params"),
-                    params={"quality": quality, "desired_fps": desired_fps},
+                    url=self._make_url("atx/power"),
+                    params={"action": action},
                 ) as response:
-                    response.raise_for_status()
-        except aiohttp.ClientError as err:
-            raise KvmdError(err)
+                    aiotools.raise_not_200(response)
+                    return True
+        except aiohttp.ClientResponseError as err:
+            if err.status == 409:
+                return False
+            raise
 
 
 # =====
@@ -112,6 +119,7 @@ class KvmdClient(_BaseClientPart):
         port: int,
         unix_path: str,
         timeout: float,
+        user_agent: str,
     ) -> None:
 
         kwargs: Dict = {
@@ -119,18 +127,17 @@ class KvmdClient(_BaseClientPart):
             "port": port,
             "unix_path": unix_path,
             "timeout": timeout,
+            "user_agent": user_agent,
         }
 
         super().__init__(**kwargs)
 
         self.auth = _AuthClientPart(**kwargs)
         self.streamer = _StreamerClientPart(**kwargs)
+        self.atx = _AtxClientPart(**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)
+        async with self._make_session(user, passwd) as session:
+            async with session.ws_connect(self._make_url("ws")) as ws:
+                yield ws