mirror of
https://github.com/mofeng-git/One-KVM.git
synced 2025-12-16 19:30:23 +08:00
288 lines
10 KiB
Python
288 lines
10 KiB
Python
# ========================================================================== #
|
|
# #
|
|
# KVMD - The main PiKVM daemon. #
|
|
# #
|
|
# Copyright (C) 2020 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 io
|
|
import contextlib
|
|
import dataclasses
|
|
import functools
|
|
import types
|
|
|
|
from typing import Callable
|
|
from typing import Awaitable
|
|
from typing import Generator
|
|
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):
|
|
pass
|
|
|
|
|
|
class StreamerTempError(StreamerError):
|
|
pass
|
|
|
|
|
|
class StreamerPermError(StreamerError):
|
|
pass
|
|
|
|
|
|
# =====
|
|
class StreamerFormats:
|
|
JPEG = 1195724874 # V4L2_PIX_FMT_JPEG
|
|
H264 = 875967048 # V4L2_PIX_FMT_H264
|
|
_MJPEG = 1196444237 # V4L2_PIX_FMT_MJPEG
|
|
|
|
|
|
class BaseStreamerClient:
|
|
def get_format(self) -> int:
|
|
raise NotImplementedError()
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def reading(self) -> AsyncGenerator[Callable[[bool], Awaitable[dict]], None]:
|
|
if self is not None: # XXX: Vulture and pylint hack
|
|
raise NotImplementedError()
|
|
yield
|
|
|
|
|
|
# =====
|
|
@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_reading_handle_errors() -> Generator[None, None, None]:
|
|
try:
|
|
yield
|
|
except Exception as ex: # Тут бывают и ассерты, и KeyError, и прочая херня
|
|
if isinstance(ex, StreamerTempError):
|
|
raise
|
|
raise StreamerTempError(tools.efmt(ex))
|
|
|
|
|
|
class HttpStreamerClient(BaseHttpClient, BaseStreamerClient):
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
unix_path: str,
|
|
timeout: float,
|
|
user_agent: str,
|
|
) -> None:
|
|
|
|
super().__init__(unix_path, timeout, user_agent)
|
|
self.__name = name
|
|
|
|
def make_session(self) -> HttpStreamerClientSession:
|
|
return HttpStreamerClientSession(self._make_http_session)
|
|
|
|
def get_format(self) -> int:
|
|
return StreamerFormats.JPEG
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def reading(self) -> AsyncGenerator[Callable[[bool], Awaitable[dict]], None]:
|
|
with _http_reading_handle_errors():
|
|
async with self._make_http_session() as session:
|
|
async with session.get(
|
|
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_reading_handle_errors():
|
|
frame = await reader.next() # pylint: disable=not-callable
|
|
if not isinstance(frame, aiohttp.BodyPartReader):
|
|
raise StreamerTempError("Expected body part")
|
|
|
|
data = bytes(await frame.read())
|
|
if not data:
|
|
raise StreamerTempError("Reached EOF")
|
|
|
|
return {
|
|
"online": (frame.headers["X-UStreamer-Online"] == "true"),
|
|
"width": int(frame.headers["X-UStreamer-Width"]),
|
|
"height": int(frame.headers["X-UStreamer-Height"]),
|
|
"data": data,
|
|
"format": StreamerFormats.JPEG,
|
|
}
|
|
|
|
yield read_frame
|
|
|
|
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.
|
|
|
|
orig_read = reader.read
|
|
|
|
async def read(self: aiohttp.StreamReader, n: int=-1) -> bytes: # pylint: disable=invalid-name
|
|
if self.is_eof():
|
|
raise StreamerTempError("StreamReader.read(): Reached EOF")
|
|
return (await orig_read(n))
|
|
|
|
reader.read = types.MethodType(read, reader) # type: ignore
|
|
|
|
def __str__(self) -> str:
|
|
return f"HttpStreamerClient({self.__name})"
|
|
|
|
|
|
# =====
|
|
@contextlib.contextmanager
|
|
def _memsink_reading_handle_errors() -> Generator[None, None, None]:
|
|
try:
|
|
yield
|
|
except StreamerPermError:
|
|
raise
|
|
except FileNotFoundError as ex:
|
|
raise StreamerTempError(tools.efmt(ex))
|
|
except Exception as ex:
|
|
raise StreamerPermError(tools.efmt(ex))
|
|
|
|
|
|
class MemsinkStreamerClient(BaseStreamerClient):
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
fmt: int,
|
|
obj: str,
|
|
lock_timeout: float,
|
|
wait_timeout: float,
|
|
drop_same_frames: float,
|
|
) -> None:
|
|
|
|
self.__name = name
|
|
self.__fmt = fmt
|
|
self.__kwargs: dict = {
|
|
"obj": obj,
|
|
"lock_timeout": lock_timeout,
|
|
"wait_timeout": wait_timeout,
|
|
"drop_same_frames": drop_same_frames,
|
|
}
|
|
|
|
def get_format(self) -> int:
|
|
return self.__fmt
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def reading(self) -> AsyncGenerator[Callable[[bool], Awaitable[dict]], None]:
|
|
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 == StreamerFormats.H264)
|
|
with _memsink_reading_handle_errors():
|
|
while True:
|
|
frame = await aiotools.run_async(sink.wait_frame, key_required)
|
|
if frame is not None:
|
|
self.__check_format(frame["format"])
|
|
return frame
|
|
yield read_frame
|
|
|
|
def __check_format(self, fmt: int) -> None:
|
|
if fmt == StreamerFormats._MJPEG: # pylint: disable=protected-access
|
|
fmt = StreamerFormats.JPEG
|
|
if fmt != self.__fmt:
|
|
raise StreamerPermError("Invalid sink format")
|
|
|
|
def __str__(self) -> str:
|
|
return f"MemsinkStreamerClient({self.__name})"
|