GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2026-01-29 00:51:53 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactoring

Browse Source
This commit is contained in:
Maxim Devaev
2024-10-22 00:51:03 +03:00
parent cda32a083f
commit 0e4a70e7b9
6 changed files with 234 additions and 232 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
kvmd/clients/__init__.py
View File

@@ -18,3 +18,67 @@
# along with this program. If not, see <https://www.gnu.org/licenses/>. #
# #
# ========================================================================== #
import types
from typing import Callable
from typing import Self
import aiohttp
# =====
class BaseHttpClientSession:
def __init__(self, make_http_session: Callable[[], aiohttp.ClientSession]) -> None:
self._make_http_session = make_http_session
self.__http_session: (aiohttp.ClientSession | None) = None
def _ensure_http_session(self) -> aiohttp.ClientSession:
if not self.__http_session:
self.__http_session = self._make_http_session()
return self.__http_session
async def close(self) -> None:
if self.__http_session:
await self.__http_session.close()
self.__http_session = None
async def __aenter__(self) -> Self:
return self
async def __aexit__(
self,
_exc_type: type[BaseException],
_exc: BaseException,
_tb: types.TracebackType,
) -> None:
await self.close()
class BaseHttpClient:
def __init__(
self,
unix_path: str,
timeout: float,
user_agent: str,
) -> None:
self.__unix_path = unix_path
self.__timeout = timeout
self.__user_agent = user_agent
def make_session(self) -> BaseHttpClientSession:
raise NotImplementedError
def _make_http_session(self, headers: (dict[str, str] | None)=None) -> aiohttp.ClientSession:
return aiohttp.ClientSession(
base_url="http://localhost:0",
headers={
"User-Agent": self.__user_agent,
**(headers or {}),
},
connector=aiohttp.UnixConnector(path=self.__unix_path),
timeout=aiohttp.ClientTimeout(total=self.__timeout),
)

121
kvmd/clients/kvmd.py
View File

@@ -23,7 +23,6 @@
import asyncio
import contextlib
import struct
import types
from typing import Callable
from typing import AsyncGenerator
@@ -34,22 +33,19 @@ from .. import aiotools
from .. import htclient
from .. import htserver
from . import BaseHttpClient
from . import BaseHttpClientSession
# =====
class _BaseApiPart:
def __init__(
self,
ensure_http_session: Callable[[], aiohttp.ClientSession],
make_url: Callable[[str], str],
) -> None:
def __init__(self, ensure_http_session: Callable[[], aiohttp.ClientSession]) -> None:
self._ensure_http_session = ensure_http_session
self._make_url = make_url
async def _set_params(self, handle: str, **params: (int | str | None)) -> None:
session = self._ensure_http_session()
async with session.post(
url=self._make_url(handle),
url=handle,
params={
key: value
for (key, value) in params.items()
@@ -63,7 +59,7 @@ class _AuthApiPart(_BaseApiPart):
async def check(self) -> bool:
session = self._ensure_http_session()
try:
async with session.get(self._make_url("auth/check")) as response:
async with session.get("/auth/check") as response:
htclient.raise_not_200(response)
return True
except aiohttp.ClientResponseError as ex:
@@ -75,13 +71,13 @@ class _AuthApiPart(_BaseApiPart):
class _StreamerApiPart(_BaseApiPart):
async def get_state(self) -> dict:
session = self._ensure_http_session()
async with session.get(self._make_url("streamer")) as response:
async with session.get("/streamer") as response:
htclient.raise_not_200(response)
return (await response.json())["result"]
async def set_params(self, quality: (int | None)=None, desired_fps: (int | None)=None) -> None:
await self._set_params(
"streamer/set_params",
"/streamer/set_params",
quality=quality,
desired_fps=desired_fps,
)
@@ -90,7 +86,7 @@ class _StreamerApiPart(_BaseApiPart):
class _HidApiPart(_BaseApiPart):
async def get_keymaps(self) -> tuple[str, set[str]]:
session = self._ensure_http_session()
async with session.get(self._make_url("hid/keymaps")) as response:
async with session.get("/hid/keymaps") as response:
htclient.raise_not_200(response)
result = (await response.json())["result"]
return (result["keymaps"]["default"], set(result["keymaps"]["available"]))
@@ -98,7 +94,7 @@ class _HidApiPart(_BaseApiPart):
async def print(self, text: str, limit: int, keymap_name: str) -> None:
session = self._ensure_http_session()
async with session.post(
url=self._make_url("hid/print"),
url="/hid/print",
params={"limit": limit, "keymap": keymap_name},
data=text,
) as response:
@@ -106,7 +102,7 @@ class _HidApiPart(_BaseApiPart):
async def set_params(self, keyboard_output: (str | None)=None, mouse_output: (str | None)=None) -> None:
await self._set_params(
"hid/set_params",
"/hid/set_params",
keyboard_output=keyboard_output,
mouse_output=mouse_output,
)
@@ -115,7 +111,7 @@ class _HidApiPart(_BaseApiPart):
class _AtxApiPart(_BaseApiPart):
async def get_state(self) -> dict:
session = self._ensure_http_session()
async with session.get(self._make_url("atx")) as response:
async with session.get("/atx") as response:
htclient.raise_not_200(response)
return (await response.json())["result"]
@@ -123,7 +119,7 @@ class _AtxApiPart(_BaseApiPart):
session = self._ensure_http_session()
try:
async with session.post(
url=self._make_url("atx/power"),
url="/atx/power",
params={"action": action},
) as response:
htclient.raise_not_200(response)
@@ -138,7 +134,6 @@ class _AtxApiPart(_BaseApiPart):
class KvmdClientWs:
def __init__(self, ws: aiohttp.ClientWebSocketResponse) -> None:
self.__ws = ws
self.__writer_queue: "asyncio.Queue[tuple[str, dict] | bytes]" = asyncio.Queue()
self.__communicated = False
@@ -200,83 +195,25 @@ class KvmdClientWs:
await self.__writer_queue.put(struct.pack(">bbbb", 5, 0, delta_x, delta_y))
class KvmdClientSession:
def __init__(
self,
make_http_session: Callable[[], aiohttp.ClientSession],
make_url: Callable[[str], str],
) -> None:
self.__make_http_session = make_http_session
self.__make_url = make_url
self.__http_session: (aiohttp.ClientSession | None) = None
args = (self.__ensure_http_session, make_url)
self.auth = _AuthApiPart(*args)
self.streamer = _StreamerApiPart(*args)
self.hid = _HidApiPart(*args)
self.atx = _AtxApiPart(*args)
class KvmdClientSession(BaseHttpClientSession):
def __init__(self, make_http_session: Callable[[], aiohttp.ClientSession]) -> None:
super().__init__(make_http_session)
self.auth = _AuthApiPart(self._ensure_http_session)
self.streamer = _StreamerApiPart(self._ensure_http_session)
self.hid = _HidApiPart(self._ensure_http_session)
self.atx = _AtxApiPart(self._ensure_http_session)
@contextlib.asynccontextmanager
async def ws(self) -> AsyncGenerator[KvmdClientWs, None]:
session = self.__ensure_http_session()
async with session.ws_connect(self.__make_url("ws"), params={"legacy": 0}) as ws:
session = self._ensure_http_session()
async with session.ws_connect("/ws", params={"legacy": "0"}) as ws:
yield KvmdClientWs(ws)
def __ensure_http_session(self) -> aiohttp.ClientSession:
if not self.__http_session:
self.__http_session = self.__make_http_session()
return self.__http_session
async def close(self) -> None:
if self.__http_session:
await self.__http_session.close()
self.__http_session = None
async def __aenter__(self) -> "KvmdClientSession":
return self
async def __aexit__(
self,
_exc_type: type[BaseException],
_exc: BaseException,
_tb: types.TracebackType,
) -> None:
await self.close()
class KvmdClient:
def __init__(
self,
unix_path: str,
timeout: float,
user_agent: str,
) -> None:
self.__unix_path = unix_path
self.__timeout = timeout
self.__user_agent = user_agent
def make_session(self, user: str, passwd: str) -> KvmdClientSession:
return KvmdClientSession(
make_http_session=(lambda: self.__make_http_session(user, passwd)),
make_url=self.__make_url,
)
def __make_http_session(self, user: str, passwd: str) -> aiohttp.ClientSession:
return aiohttp.ClientSession(
headers={
"X-KVMD-User": user,
"X-KVMD-Passwd": passwd,
"User-Agent": self.__user_agent,
},
connector=aiohttp.UnixConnector(path=self.__unix_path),
timeout=aiohttp.ClientTimeout(total=self.__timeout),
)
def __make_url(self, handle: str) -> str:
assert not handle.startswith("/"), handle
return f"http://localhost:0/{handle}"
class KvmdClient(BaseHttpClient):
def make_session(self, user: str="", passwd: str="") -> KvmdClientSession:
headers = {
"X-KVMD-User": user,
"X-KVMD-Passwd": passwd,
}
return KvmdClientSession(lambda: self._make_http_session(headers))

133
kvmd/clients/streamer.py
View File

@@ -20,7 +20,10 @@
# ========================================================================== #
import io
import contextlib
import dataclasses
import functools
import types
from typing import Callable
@@ -31,10 +34,15 @@ from typing import AsyncGenerator
import aiohttp
import ustreamer
from PIL import Image as PilImage
from .. import tools
from .. import aiotools
from .. import htclient
from . import BaseHttpClient
from . import BaseHttpClientSession
# =====
class StreamerError(Exception):
@@ -50,7 +58,7 @@ class StreamerPermError(StreamerError):
# =====
class StreamFormats:
class StreamerFormats:
JPEG = 1195724874 # V4L2_PIX_FMT_JPEG
H264 = 875967048 # V4L2_PIX_FMT_H264
_MJPEG = 1196444237 # V4L2_PIX_FMT_MJPEG
@@ -68,8 +76,76 @@ class BaseStreamerClient:
# =====
@dataclasses.dataclass(frozen=True)
class StreamerSnapshot:
online: bool
width: int
height: int
headers: tuple[tuple[str, str], ...]
data: bytes
async def make_preview(self, max_width: int, max_height: int, quality: int) -> bytes:
assert max_width >= 0
assert max_height >= 0
assert quality > 0
if max_width == 0 and max_height == 0:
max_width = self.width // 5
max_height = self.height // 5
else:
max_width = min((max_width or self.width), self.width)
max_height = min((max_height or self.height), self.height)
if (max_width, max_height) == (self.width, self.height):
return self.data
return (await aiotools.run_async(self.__inner_make_preview, max_width, max_height, quality))
@functools.lru_cache(maxsize=1)
def __inner_make_preview(self, max_width: int, max_height: int, quality: int) -> bytes:
with io.BytesIO(self.data) as snapshot_bio:
with io.BytesIO() as preview_bio:
with PilImage.open(snapshot_bio) as image:
image.thumbnail((max_width, max_height), PilImage.Resampling.LANCZOS)
image.save(preview_bio, format="jpeg", quality=quality)
return preview_bio.getvalue()
class HttpStreamerClientSession(BaseHttpClientSession):
async def get_state(self) -> dict:
session = self._ensure_http_session()
async with session.get("/state") as response:
htclient.raise_not_200(response)
return (await response.json())["result"]
async def take_snapshot(self, timeout: float) -> StreamerSnapshot:
session = self._ensure_http_session()
async with session.get(
url="/snapshot",
timeout=aiohttp.ClientTimeout(total=timeout),
) as response:
htclient.raise_not_200(response)
return StreamerSnapshot(
online=(response.headers["X-UStreamer-Online"] == "true"),
width=int(response.headers["X-UStreamer-Width"]),
height=int(response.headers["X-UStreamer-Height"]),
headers=tuple(
(key, value)
for (key, value) in tools.sorted_kvs(dict(response.headers))
if key.lower().startswith("x-ustreamer-") or key.lower() in [
"x-timestamp",
"access-control-allow-origin",
"cache-control",
"pragma",
"expires",
]
),
data=bytes(await response.read()),
)
@contextlib.contextmanager
def _http_handle_errors() -> Generator[None, None, None]:
def _http_reading_handle_errors() -> Generator[None, None, None]:
try:
yield
except Exception as ex: # Тут бывают и ассерты, и KeyError, и прочая херня
@@ -78,7 +154,7 @@ def _http_handle_errors() -> Generator[None, None, None]:
raise StreamerTempError(tools.efmt(ex))
class HttpStreamerClient(BaseStreamerClient):
class HttpStreamerClient(BaseHttpClient, BaseStreamerClient):
def __init__(
self,
name: str,
@@ -87,29 +163,35 @@ class HttpStreamerClient(BaseStreamerClient):
user_agent: str,
) -> None:
super().__init__(unix_path, timeout, user_agent)
self.__name = name
self.__unix_path = unix_path
self.__timeout = timeout
self.__user_agent = user_agent
def make_session(self) -> HttpStreamerClientSession:
return HttpStreamerClientSession(self._make_http_session)
def get_format(self) -> int:
return StreamFormats.JPEG
return StreamerFormats.JPEG
@contextlib.asynccontextmanager
async def reading(self) -> AsyncGenerator[Callable[[bool], Awaitable[dict]], None]:
with _http_handle_errors():
async with self.__make_http_session() as session:
with _http_reading_handle_errors():
async with self._make_http_session() as session:
async with session.get(
url=self.__make_url("stream"),
url="/stream",
params={"extra_headers": "1"},
timeout=aiohttp.ClientTimeout(
connect=session.timeout.total,
sock_read=session.timeout.total,
),
) as response:
htclient.raise_not_200(response)
reader = aiohttp.MultipartReader.from_response(response)
self.__patch_stream_reader(reader.resp.content)
async def read_frame(key_required: bool) -> dict:
_ = key_required
with _http_handle_errors():
with _http_reading_handle_errors():
frame = await reader.next() # pylint: disable=not-callable
if not isinstance(frame, aiohttp.BodyPartReader):
raise StreamerTempError("Expected body part")
@@ -123,26 +205,11 @@ class HttpStreamerClient(BaseStreamerClient):
"width": int(frame.headers["X-UStreamer-Width"]),
"height": int(frame.headers["X-UStreamer-Height"]),
"data": data,
"format": StreamFormats.JPEG,
"format": StreamerFormats.JPEG,
}
yield read_frame
def __make_http_session(self) -> aiohttp.ClientSession:
kwargs: dict = {
"headers": {"User-Agent": self.__user_agent},
"connector": aiohttp.UnixConnector(path=self.__unix_path),
"timeout": aiohttp.ClientTimeout(
connect=self.__timeout,
sock_read=self.__timeout,
),
}
return aiohttp.ClientSession(**kwargs)
def __make_url(self, handle: str) -> str:
assert not handle.startswith("/"), handle
return f"http://localhost:0/{handle}"
def __patch_stream_reader(self, reader: aiohttp.StreamReader) -> None:
# https://github.com/pikvm/pikvm/issues/92
# Infinite looping in BodyPartReader.read() because _at_eof flag.
@@ -162,7 +229,7 @@ class HttpStreamerClient(BaseStreamerClient):
# =====
@contextlib.contextmanager
def _memsink_handle_errors() -> Generator[None, None, None]:
def _memsink_reading_handle_errors() -> Generator[None, None, None]:
try:
yield
except StreamerPermError:
@@ -198,11 +265,11 @@ class MemsinkStreamerClient(BaseStreamerClient):
@contextlib.asynccontextmanager
async def reading(self) -> AsyncGenerator[Callable[[bool], Awaitable[dict]], None]:
with _memsink_handle_errors():
with _memsink_reading_handle_errors():
with ustreamer.Memsink(**self.__kwargs) as sink:
async def read_frame(key_required: bool) -> dict:
key_required = (key_required and self.__fmt == StreamFormats.H264)
with _memsink_handle_errors():
key_required = (key_required and self.__fmt == StreamerFormats.H264)
with _memsink_reading_handle_errors():
while True:
frame = await aiotools.run_async(sink.wait_frame, key_required)
if frame is not None:
@@ -211,8 +278,8 @@ class MemsinkStreamerClient(BaseStreamerClient):
yield read_frame
def __check_format(self, fmt: int) -> None:
if fmt == StreamFormats._MJPEG: # pylint: disable=protected-access
fmt = StreamFormats.JPEG
if fmt == StreamerFormats._MJPEG: # pylint: disable=protected-access
fmt = StreamerFormats.JPEG
if fmt != self.__fmt:
raise StreamerPermError("Invalid sink format")