refactoring

This commit is contained in:
Devaev Maxim 2020-05-28 11:03:49 +03:00
parent fbdfb009a1
commit 1c93f6a562
9 changed files with 63 additions and 42 deletions

View File

@ -21,8 +21,3 @@
__version__ = "1.63" __version__ = "1.63"
# =====
def make_user_agent(app: str) -> str:
return f"{app}/{__version__}"

View File

@ -39,8 +39,6 @@ from typing import TypeVar
from typing import Optional from typing import Optional
from typing import Any from typing import Any
import aiohttp
import aiofiles import aiofiles
import aiofiles.base import aiofiles.base
@ -95,20 +93,6 @@ async def wait_first(*aws: Awaitable) -> Tuple[Set[asyncio.Future], Set[asyncio.
return (await asyncio.wait(list(aws), return_when=asyncio.FIRST_COMPLETED)) return (await asyncio.wait(list(aws), return_when=asyncio.FIRST_COMPLETED))
# =====
def raise_not_200(response: aiohttp.ClientResponse) -> None:
if response.status != 200:
assert response.reason is not None
response.release()
raise aiohttp.ClientResponseError(
response.request_info,
response.history,
status=response.status,
message=response.reason,
headers=response.headers,
)
# ===== # =====
async def afile_write_now(afile: aiofiles.base.AiofilesContextManager, data: bytes) -> None: async def afile_write_now(afile: aiofiles.base.AiofilesContextManager, data: bytes) -> None:
await afile.write(data) await afile.write(data)

View File

@ -25,7 +25,7 @@ from typing import Optional
from ...clients.kvmd import KvmdClient from ...clients.kvmd import KvmdClient
from ... import make_user_agent from ... import htclient
from .. import init from .. import init
@ -45,7 +45,7 @@ def main(argv: Optional[List[str]]=None) -> None:
IpmiServer( IpmiServer(
auth_manager=IpmiAuthManager(**config.auth._unpack()), auth_manager=IpmiAuthManager(**config.auth._unpack()),
kvmd=KvmdClient( kvmd=KvmdClient(
user_agent=make_user_agent("KVMD-IPMI"), user_agent=htclient.make_user_agent("KVMD-IPMI"),
**config.kvmd._unpack(), **config.kvmd._unpack(),
), ),
**config.server._unpack(), **config.server._unpack(),

View File

@ -36,10 +36,9 @@ import aiohttp
from ...logging import get_logger from ...logging import get_logger
from ... import aiotools from ... import aiotools
from ... import htclient
from ... import gpio from ... import gpio
from ... import make_user_agent
# ===== # =====
class Streamer: # pylint: disable=too-many-instance-attributes class Streamer: # pylint: disable=too-many-instance-attributes
@ -182,7 +181,7 @@ class Streamer: # pylint: disable=too-many-instance-attributes
session = self.__ensure_http_session() session = self.__ensure_http_session()
try: try:
async with session.get(self.__make_url("state")) as response: async with session.get(self.__make_url("state")) as response:
aiotools.raise_not_200(response) htclient.raise_not_200(response)
state = (await response.json())["result"] state = (await response.json())["result"]
except (aiohttp.ClientConnectionError, aiohttp.ServerConnectionError): except (aiohttp.ClientConnectionError, aiohttp.ServerConnectionError):
pass pass
@ -245,7 +244,7 @@ class Streamer: # pylint: disable=too-many-instance-attributes
def __ensure_http_session(self) -> aiohttp.ClientSession: def __ensure_http_session(self) -> aiohttp.ClientSession:
if not self.__http_session: if not self.__http_session:
kwargs: Dict = { kwargs: Dict = {
"headers": {"User-Agent": make_user_agent("KVMD")}, "headers": {"User-Agent": htclient.make_user_agent("KVMD")},
"timeout": aiohttp.ClientTimeout(total=self.__timeout), "timeout": aiohttp.ClientTimeout(total=self.__timeout),
} }
if self.__unix_path: if self.__unix_path:

View File

@ -26,7 +26,7 @@ from typing import Optional
from ...clients.kvmd import KvmdClient from ...clients.kvmd import KvmdClient
from ...clients.streamer import StreamerClient from ...clients.streamer import StreamerClient
from ... import make_user_agent from ... import htclient
from .. import init from .. import init
@ -42,7 +42,7 @@ def main(argv: Optional[List[str]]=None) -> None:
argv=argv, argv=argv,
)[2].vnc )[2].vnc
user_agent = make_user_agent("KVMD-VNC") user_agent = htclient.make_user_agent("KVMD-VNC")
# pylint: disable=protected-access # pylint: disable=protected-access
VncServer( VncServer(

View File

@ -37,6 +37,7 @@ from typing import Optional
import aiohttp import aiohttp
from .. import aiotools from .. import aiotools
from .. import htclient
# ===== # =====
@ -56,7 +57,7 @@ class _AuthApiPart(_BaseApiPart):
session = self._ensure_http_session() session = self._ensure_http_session()
try: try:
async with session.get(self._make_url("auth/check")) as response: async with session.get(self._make_url("auth/check")) as response:
aiotools.raise_not_200(response) htclient.raise_not_200(response)
return True return True
except aiohttp.ClientResponseError as err: except aiohttp.ClientResponseError as err:
if err.status in [401, 403]: if err.status in [401, 403]:
@ -71,14 +72,14 @@ class _StreamerApiPart(_BaseApiPart):
url=self._make_url("streamer/set_params"), url=self._make_url("streamer/set_params"),
params={"quality": quality, "desired_fps": desired_fps}, params={"quality": quality, "desired_fps": desired_fps},
) as response: ) as response:
aiotools.raise_not_200(response) htclient.raise_not_200(response)
class _HidApiPart(_BaseApiPart): class _HidApiPart(_BaseApiPart):
async def get_keymaps(self) -> Tuple[str, Set[str]]: async def get_keymaps(self) -> Tuple[str, Set[str]]:
session = self._ensure_http_session() session = self._ensure_http_session()
async with session.get(self._make_url("hid/keymaps")) as response: async with session.get(self._make_url("hid/keymaps")) as response:
aiotools.raise_not_200(response) htclient.raise_not_200(response)
result = (await response.json())["result"] result = (await response.json())["result"]
return (result["keymaps"]["default"], set(result["keymaps"]["available"])) return (result["keymaps"]["default"], set(result["keymaps"]["available"]))
@ -89,14 +90,14 @@ class _HidApiPart(_BaseApiPart):
params={"limit": limit, "keymap": keymap_name}, params={"limit": limit, "keymap": keymap_name},
data=text, data=text,
) as response: ) as response:
aiotools.raise_not_200(response) htclient.raise_not_200(response)
class _AtxApiPart(_BaseApiPart): class _AtxApiPart(_BaseApiPart):
async def get_state(self) -> Dict: async def get_state(self) -> Dict:
session = self._ensure_http_session() session = self._ensure_http_session()
async with session.get(self._make_url("atx")) as response: async with session.get(self._make_url("atx")) as response:
aiotools.raise_not_200(response) htclient.raise_not_200(response)
return (await response.json())["result"] return (await response.json())["result"]
async def switch_power(self, action: str) -> bool: async def switch_power(self, action: str) -> bool:
@ -106,7 +107,7 @@ class _AtxApiPart(_BaseApiPart):
url=self._make_url("atx/power"), url=self._make_url("atx/power"),
params={"action": action}, params={"action": action},
) as response: ) as response:
aiotools.raise_not_200(response) htclient.raise_not_200(response)
return True return True
except aiohttp.ClientResponseError as err: except aiohttp.ClientResponseError as err:
if err.status == 409: if err.status == 409:

View File

@ -26,7 +26,7 @@ from typing import AsyncGenerator
import aiohttp import aiohttp
from .. import aiotools from .. import htclient
# ===== # =====
@ -59,7 +59,7 @@ class StreamerClient:
url=self.__make_url("stream"), url=self.__make_url("stream"),
params={"extra_headers": "1"}, params={"extra_headers": "1"},
) as response: ) as response:
aiotools.raise_not_200(response) htclient.raise_not_200(response)
reader = aiohttp.MultipartReader.from_response(response) reader = aiohttp.MultipartReader.from_response(response)
while True: while True:
@ -87,7 +87,7 @@ class StreamerClient:
# async def get_snapshot(self) -> Tuple[bool, bytes]: # async def get_snapshot(self) -> Tuple[bool, bytes]:
# async with self.__make_http_session(infinite=False) as session: # async with self.__make_http_session(infinite=False) as session:
# async with session.get(self.__make_url("snapshot")) as response: # async with session.get(self.__make_url("snapshot")) as response:
# aiotools.raise_not_200(response) # htclient.raise_not_200(response)
# return ( # return (
# (response.headers["X-UStreamer-Online"] == "true"), # (response.headers["X-UStreamer-Online"] == "true"),
# bytes(await response.read()), # bytes(await response.read()),

43
kvmd/htclient.py Normal file
View File

@ -0,0 +1,43 @@
# ========================================================================== #
# #
# KVMD - The main Pi-KVM daemon. #
# #
# Copyright (C) 2018 Maxim Devaev <mdevaev@gmail.com> #
# #
# This program is free software: you can redistribute it and/or modify #
# it under the terms of the GNU General Public License as published by #
# the Free Software Foundation, either version 3 of the License, or #
# (at your option) any later version. #
# #
# This program is distributed in the hope that it will be useful, #
# but WITHOUT ANY WARRANTY; without even the implied warranty of #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the #
# GNU General Public License for more details. #
# #
# You should have received a copy of the GNU General Public License #
# along with this program. If not, see <https://www.gnu.org/licenses/>. #
# #
# ========================================================================== #
import aiohttp
from . import __version__
# =====
def make_user_agent(app: str) -> str:
return f"{app}/{__version__}"
def raise_not_200(response: aiohttp.ClientResponse) -> None:
if response.status != 200:
assert response.reason is not None
response.release()
raise aiohttp.ClientResponseError(
response.request_info,
response.history,
status=response.status,
message=response.reason,
headers=response.headers,
)

View File

@ -33,8 +33,7 @@ from ...validators.basic import valid_float_f01
from ...logging import get_logger from ...logging import get_logger
from ... import make_user_agent from ... import htclient
from ... import aiotools
from . import BaseAuthService from . import BaseAuthService
@ -86,11 +85,11 @@ class Plugin(BaseAuthService):
"secret": self.__secret, "secret": self.__secret,
}, },
headers={ headers={
"User-Agent": make_user_agent("KVMD"), "User-Agent": htclient.make_user_agent("KVMD"),
"X-KVMD-User": user, "X-KVMD-User": user,
}, },
) as response: ) as response:
aiotools.raise_not_200(response) htclient.raise_not_200(response)
return True return True
except Exception: except Exception:
get_logger().exception("Failed HTTP auth request for user %r", user) get_logger().exception("Failed HTTP auth request for user %r", user)