mirror of
https://github.com/mofeng-git/One-KVM.git
synced 2025-12-13 01:30:31 +08:00
449 lines
16 KiB
Python
449 lines
16 KiB
Python
# ========================================================================== #
|
|
# #
|
|
# KVMD - The main PiKVM daemon. #
|
|
# #
|
|
# Copyright (C) 2018-2024 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 signal
|
|
import asyncio
|
|
import asyncio.subprocess
|
|
import dataclasses
|
|
import copy
|
|
|
|
from typing import AsyncGenerator
|
|
from typing import Any
|
|
|
|
import aiohttp
|
|
|
|
from ...logging import get_logger
|
|
|
|
from ...clients.streamer import StreamerSnapshot
|
|
from ...clients.streamer import HttpStreamerClient
|
|
from ...clients.streamer import HttpStreamerClientSession
|
|
|
|
from ... import tools
|
|
from ... import aiotools
|
|
from ... import aioproc
|
|
from ... import htclient
|
|
|
|
|
|
# =====
|
|
class _StreamerParams:
|
|
__DESIRED_FPS = "desired_fps"
|
|
|
|
__QUALITY = "quality"
|
|
|
|
__RESOLUTION = "resolution"
|
|
__AVAILABLE_RESOLUTIONS = "available_resolutions"
|
|
|
|
__H264_BITRATE = "h264_bitrate"
|
|
__H264_GOP = "h264_gop"
|
|
|
|
def __init__( # pylint: disable=too-many-arguments
|
|
self,
|
|
quality: int,
|
|
|
|
resolution: str,
|
|
available_resolutions: list[str],
|
|
|
|
desired_fps: int,
|
|
desired_fps_min: int,
|
|
desired_fps_max: int,
|
|
|
|
h264_bitrate: int,
|
|
h264_bitrate_min: int,
|
|
h264_bitrate_max: int,
|
|
|
|
h264_gop: int,
|
|
h264_gop_min: int,
|
|
h264_gop_max: int,
|
|
) -> None:
|
|
|
|
self.__has_quality = bool(quality)
|
|
self.__has_resolution = bool(resolution)
|
|
self.__has_h264 = bool(h264_bitrate)
|
|
|
|
self.__params: dict = {self.__DESIRED_FPS: min(max(desired_fps, desired_fps_min), desired_fps_max)}
|
|
self.__limits: dict = {self.__DESIRED_FPS: {"min": desired_fps_min, "max": desired_fps_max}}
|
|
|
|
if self.__has_quality:
|
|
self.__params[self.__QUALITY] = quality
|
|
|
|
if self.__has_resolution:
|
|
self.__params[self.__RESOLUTION] = resolution
|
|
self.__limits[self.__AVAILABLE_RESOLUTIONS] = available_resolutions
|
|
|
|
if self.__has_h264:
|
|
self.__params[self.__H264_BITRATE] = min(max(h264_bitrate, h264_bitrate_min), h264_bitrate_max)
|
|
self.__limits[self.__H264_BITRATE] = {"min": h264_bitrate_min, "max": h264_bitrate_max}
|
|
self.__params[self.__H264_GOP] = min(max(h264_gop, h264_gop_min), h264_gop_max)
|
|
self.__limits[self.__H264_GOP] = {"min": h264_gop_min, "max": h264_gop_max}
|
|
|
|
def get_features(self) -> dict:
|
|
return {
|
|
self.__QUALITY: self.__has_quality,
|
|
self.__RESOLUTION: self.__has_resolution,
|
|
"h264": self.__has_h264,
|
|
}
|
|
|
|
def get_limits(self) -> dict:
|
|
limits = copy.deepcopy(self.__limits)
|
|
if self.__has_resolution:
|
|
limits[self.__AVAILABLE_RESOLUTIONS] = list(limits[self.__AVAILABLE_RESOLUTIONS])
|
|
return limits
|
|
|
|
def get_params(self) -> dict:
|
|
return dict(self.__params)
|
|
|
|
def set_params(self, params: dict) -> None:
|
|
new_params = dict(self.__params)
|
|
|
|
if self.__QUALITY in params and self.__has_quality:
|
|
new_params[self.__QUALITY] = min(max(params[self.__QUALITY], 1), 100)
|
|
|
|
if self.__RESOLUTION in params and self.__has_resolution:
|
|
if params[self.__RESOLUTION] in self.__limits[self.__AVAILABLE_RESOLUTIONS]:
|
|
new_params[self.__RESOLUTION] = params[self.__RESOLUTION]
|
|
|
|
for (key, enabled) in [
|
|
(self.__DESIRED_FPS, True),
|
|
(self.__H264_BITRATE, self.__has_h264),
|
|
(self.__H264_GOP, self.__has_h264),
|
|
]:
|
|
if key in params and enabled:
|
|
if self.__check_limits_min_max(key, params[key]):
|
|
new_params[key] = params[key]
|
|
|
|
self.__params = new_params
|
|
|
|
def __check_limits_min_max(self, key: str, value: int) -> bool:
|
|
return (self.__limits[key]["min"] <= value <= self.__limits[key]["max"])
|
|
|
|
|
|
class Streamer: # pylint: disable=too-many-instance-attributes
|
|
__ST_FULL = 0xFF
|
|
__ST_PARAMS = 0x01
|
|
__ST_STREAMER = 0x02
|
|
__ST_SNAPSHOT = 0x04
|
|
|
|
def __init__( # pylint: disable=too-many-arguments,too-many-locals
|
|
self,
|
|
|
|
reset_delay: float,
|
|
shutdown_delay: float,
|
|
state_poll: float,
|
|
|
|
unix_path: str,
|
|
timeout: float,
|
|
snapshot_timeout: float,
|
|
|
|
process_name_prefix: str,
|
|
|
|
pre_start_cmd: list[str],
|
|
pre_start_cmd_remove: list[str],
|
|
pre_start_cmd_append: list[str],
|
|
|
|
cmd: list[str],
|
|
cmd_remove: list[str],
|
|
cmd_append: list[str],
|
|
|
|
post_stop_cmd: list[str],
|
|
post_stop_cmd_remove: list[str],
|
|
post_stop_cmd_append: list[str],
|
|
|
|
**params_kwargs: Any,
|
|
) -> None:
|
|
|
|
self.__reset_delay = reset_delay
|
|
self.__shutdown_delay = shutdown_delay
|
|
self.__state_poll = state_poll
|
|
|
|
self.__unix_path = unix_path
|
|
self.__snapshot_timeout = snapshot_timeout
|
|
|
|
self.__process_name_prefix = process_name_prefix
|
|
|
|
self.__pre_start_cmd = tools.build_cmd(pre_start_cmd, pre_start_cmd_remove, pre_start_cmd_append)
|
|
self.__cmd = tools.build_cmd(cmd, cmd_remove, cmd_append)
|
|
self.__post_stop_cmd = tools.build_cmd(post_stop_cmd, post_stop_cmd_remove, post_stop_cmd_append)
|
|
|
|
self.__params = _StreamerParams(**params_kwargs)
|
|
|
|
self.__stop_task: (asyncio.Task | None) = None
|
|
self.__stop_wip = False
|
|
|
|
self.__streamer_task: (asyncio.Task | None) = None
|
|
self.__streamer_proc: (asyncio.subprocess.Process | None) = None # pylint: disable=no-member
|
|
|
|
self.__client = HttpStreamerClient(
|
|
name="jpeg",
|
|
unix_path=self.__unix_path,
|
|
timeout=timeout,
|
|
user_agent=htclient.make_user_agent("KVMD"),
|
|
)
|
|
self.__client_session: (HttpStreamerClientSession | None) = None
|
|
|
|
self.__snapshot: (StreamerSnapshot | None) = None
|
|
|
|
self.__notifier = aiotools.AioNotifier()
|
|
|
|
# =====
|
|
|
|
@aiotools.atomic_fg
|
|
async def ensure_start(self, reset: bool) -> None:
|
|
if not self.__streamer_task or self.__stop_task:
|
|
logger = get_logger(0)
|
|
|
|
if self.__stop_task:
|
|
if not self.__stop_wip:
|
|
self.__stop_task.cancel()
|
|
await asyncio.gather(self.__stop_task, return_exceptions=True)
|
|
logger.info("Streamer stop cancelled")
|
|
return
|
|
else:
|
|
await asyncio.gather(self.__stop_task, return_exceptions=True)
|
|
|
|
if reset and self.__reset_delay > 0:
|
|
logger.info("Waiting %.2f seconds for reset delay ...", self.__reset_delay)
|
|
await asyncio.sleep(self.__reset_delay)
|
|
logger.info("Starting streamer ...")
|
|
await self.__inner_start()
|
|
|
|
@aiotools.atomic_fg
|
|
async def ensure_stop(self, immediately: bool) -> None:
|
|
if self.__streamer_task:
|
|
logger = get_logger(0)
|
|
|
|
if immediately:
|
|
if self.__stop_task:
|
|
if not self.__stop_wip:
|
|
self.__stop_task.cancel()
|
|
await asyncio.gather(self.__stop_task, return_exceptions=True)
|
|
logger.info("Stopping streamer immediately ...")
|
|
await self.__inner_stop()
|
|
else:
|
|
await asyncio.gather(self.__stop_task, return_exceptions=True)
|
|
else:
|
|
logger.info("Stopping streamer immediately ...")
|
|
await self.__inner_stop()
|
|
|
|
elif not self.__stop_task:
|
|
|
|
async def delayed_stop() -> None:
|
|
try:
|
|
await asyncio.sleep(self.__shutdown_delay)
|
|
self.__stop_wip = True
|
|
logger.info("Stopping streamer after delay ...")
|
|
await self.__inner_stop()
|
|
finally:
|
|
self.__stop_task = None
|
|
self.__stop_wip = False
|
|
|
|
logger.info("Planning to stop streamer in %.2f seconds ...", self.__shutdown_delay)
|
|
self.__stop_task = asyncio.create_task(delayed_stop())
|
|
|
|
def is_working(self) -> bool:
|
|
# Запущено и не планирует останавливаться
|
|
return bool(self.__streamer_task and not self.__stop_task)
|
|
|
|
# =====
|
|
|
|
def set_params(self, params: dict) -> None:
|
|
assert not self.__streamer_task
|
|
self.__notifier.notify(self.__ST_PARAMS)
|
|
return self.__params.set_params(params)
|
|
|
|
def get_params(self) -> dict:
|
|
return self.__params.get_params()
|
|
|
|
# =====
|
|
|
|
async def get_state(self) -> dict:
|
|
return {
|
|
"features": self.__params.get_features(),
|
|
"limits": self.__params.get_limits(),
|
|
"params": self.__params.get_params(),
|
|
"streamer": (await self.__get_streamer_state()),
|
|
"snapshot": self.__get_snapshot_state(),
|
|
}
|
|
|
|
async def trigger_state(self) -> None:
|
|
self.__notifier.notify(self.__ST_FULL)
|
|
|
|
async def poll_state(self) -> AsyncGenerator[dict, None]:
|
|
def signal_handler(*_: Any) -> None:
|
|
get_logger(0).info("Got SIGUSR2, checking the stream state ...")
|
|
self.__notifier.notify(self.__ST_STREAMER)
|
|
|
|
get_logger(0).info("Installing SIGUSR2 streamer handler ...")
|
|
asyncio.get_event_loop().add_signal_handler(signal.SIGUSR2, signal_handler)
|
|
|
|
prev: dict = {}
|
|
while True:
|
|
new: dict = {}
|
|
|
|
mask = await self.__notifier.wait(timeout=self.__state_poll)
|
|
if mask == self.__ST_FULL:
|
|
new = await self.get_state()
|
|
prev = copy.deepcopy(new)
|
|
yield new
|
|
continue
|
|
|
|
if mask < 0:
|
|
mask = self.__ST_STREAMER
|
|
|
|
def check_update(key: str, value: (dict | None)) -> None:
|
|
if prev.get(key) != value:
|
|
new[key] = value
|
|
|
|
if mask & self.__ST_PARAMS:
|
|
check_update("params", self.__params.get_params())
|
|
if mask & self.__ST_STREAMER:
|
|
check_update("streamer", await self.__get_streamer_state())
|
|
if mask & self.__ST_SNAPSHOT:
|
|
check_update("snapshot", self.__get_snapshot_state())
|
|
|
|
if new and prev != new:
|
|
prev.update(copy.deepcopy(new))
|
|
yield new
|
|
|
|
async def __get_streamer_state(self) -> (dict | None):
|
|
if self.__streamer_task:
|
|
session = self.__ensure_client_session()
|
|
try:
|
|
return (await session.get_state())
|
|
except (aiohttp.ClientConnectionError, aiohttp.ServerConnectionError):
|
|
pass
|
|
except Exception:
|
|
get_logger().exception("Invalid streamer response from /state")
|
|
return None
|
|
|
|
def __get_snapshot_state(self) -> dict:
|
|
if self.__snapshot:
|
|
snapshot = dataclasses.asdict(self.__snapshot)
|
|
del snapshot["headers"]
|
|
del snapshot["data"]
|
|
return {"saved": snapshot}
|
|
return {"saved": None}
|
|
|
|
# =====
|
|
|
|
async def take_snapshot(self, save: bool, load: bool, allow_offline: bool) -> (StreamerSnapshot | None):
|
|
if load:
|
|
return self.__snapshot
|
|
logger = get_logger()
|
|
session = self.__ensure_client_session()
|
|
try:
|
|
snapshot = await session.take_snapshot(self.__snapshot_timeout)
|
|
if snapshot.online or allow_offline:
|
|
if save:
|
|
self.__snapshot = snapshot
|
|
self.__notifier.notify(self.__ST_SNAPSHOT)
|
|
return snapshot
|
|
logger.error("Stream is offline, no signal or so")
|
|
except (aiohttp.ClientConnectionError, aiohttp.ServerConnectionError) as ex:
|
|
logger.error("Can't connect to streamer: %s", tools.efmt(ex))
|
|
except Exception:
|
|
logger.exception("Invalid streamer response from /snapshot")
|
|
return None
|
|
|
|
def remove_snapshot(self) -> None:
|
|
self.__snapshot = None
|
|
|
|
# =====
|
|
|
|
@aiotools.atomic_fg
|
|
async def cleanup(self) -> None:
|
|
await self.ensure_stop(immediately=True)
|
|
if self.__client_session:
|
|
await self.__client_session.close()
|
|
self.__client_session = None
|
|
|
|
def __ensure_client_session(self) -> HttpStreamerClientSession:
|
|
if not self.__client_session:
|
|
self.__client_session = self.__client.make_session()
|
|
return self.__client_session
|
|
|
|
# =====
|
|
|
|
@aiotools.atomic_fg
|
|
async def __inner_start(self) -> None:
|
|
assert not self.__streamer_task
|
|
await self.__run_hook("PRE-START-CMD", self.__pre_start_cmd)
|
|
self.__streamer_task = asyncio.create_task(self.__streamer_task_loop())
|
|
|
|
@aiotools.atomic_fg
|
|
async def __inner_stop(self) -> None:
|
|
assert self.__streamer_task
|
|
self.__streamer_task.cancel()
|
|
await asyncio.gather(self.__streamer_task, return_exceptions=True)
|
|
await self.__kill_streamer_proc()
|
|
await self.__run_hook("POST-STOP-CMD", self.__post_stop_cmd)
|
|
self.__streamer_task = None
|
|
|
|
# =====
|
|
|
|
async def __streamer_task_loop(self) -> None: # pylint: disable=too-many-branches
|
|
logger = get_logger(0)
|
|
while True: # pylint: disable=too-many-nested-blocks
|
|
try:
|
|
await self.__start_streamer_proc()
|
|
assert self.__streamer_proc is not None
|
|
await aioproc.log_stdout_infinite(self.__streamer_proc, logger)
|
|
raise RuntimeError("Streamer unexpectedly died")
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception:
|
|
if self.__streamer_proc:
|
|
logger.exception("Unexpected streamer error: pid=%d", self.__streamer_proc.pid)
|
|
else:
|
|
logger.exception("Can't start streamer")
|
|
await self.__kill_streamer_proc()
|
|
await asyncio.sleep(1)
|
|
|
|
def __make_cmd(self, cmd: list[str]) -> list[str]:
|
|
return [
|
|
part.format(
|
|
unix=self.__unix_path,
|
|
process_name_prefix=self.__process_name_prefix,
|
|
**self.__params.get_params(),
|
|
)
|
|
for part in cmd
|
|
]
|
|
|
|
async def __run_hook(self, name: str, cmd: list[str]) -> None:
|
|
logger = get_logger()
|
|
cmd = self.__make_cmd(cmd)
|
|
logger.info("%s: %s", name, tools.cmdfmt(cmd))
|
|
try:
|
|
await aioproc.log_process(cmd, logger, prefix=name)
|
|
except Exception as ex:
|
|
logger.exception("Can't execute command: %s", ex)
|
|
|
|
async def __start_streamer_proc(self) -> None:
|
|
assert self.__streamer_proc is None
|
|
cmd = self.__make_cmd(self.__cmd)
|
|
self.__streamer_proc = await aioproc.run_process(cmd)
|
|
get_logger(0).info("Started streamer pid=%d: %s", self.__streamer_proc.pid, tools.cmdfmt(cmd))
|
|
|
|
async def __kill_streamer_proc(self) -> None:
|
|
if self.__streamer_proc:
|
|
await aioproc.kill_process(self.__streamer_proc, 1, get_logger(0))
|
|
self.__streamer_proc = None
|