mirror of
https://github.com/mofeng-git/One-KVM.git
synced 2025-12-12 01:00:29 +08:00
545 lines
22 KiB
Python
545 lines
22 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 os
|
||
import asyncio
|
||
import socket
|
||
import dataclasses
|
||
import contextlib
|
||
|
||
import aiohttp
|
||
|
||
from ...logging import get_logger
|
||
|
||
from ...keyboard.keysym import SymmapModifiers
|
||
from ...keyboard.keysym import build_symmap
|
||
from ...keyboard.mappings import WebModifiers
|
||
from ...keyboard.mappings import X11Modifiers
|
||
from ...keyboard.mappings import AT1_TO_WEB
|
||
|
||
from ...clients.kvmd import KvmdClientWs
|
||
from ...clients.kvmd import KvmdClientSession
|
||
from ...clients.kvmd import KvmdClient
|
||
|
||
from ...clients.streamer import StreamerError
|
||
from ...clients.streamer import StreamerPermError
|
||
from ...clients.streamer import StreamerFormats
|
||
from ...clients.streamer import BaseStreamerClient
|
||
|
||
from ... import tools
|
||
from ... import aiotools
|
||
from ... import network
|
||
|
||
from .rfb import RfbClient
|
||
from .rfb.stream import rfb_format_remote
|
||
from .rfb.errors import RfbError
|
||
|
||
from .vncauth import VncAuthKvmdCredentials
|
||
from .vncauth import VncAuthManager
|
||
|
||
from .render import make_text_jpeg
|
||
|
||
|
||
# =====
|
||
@dataclasses.dataclass()
|
||
class _SharedParams:
|
||
width: int = dataclasses.field(default=800)
|
||
height: int = dataclasses.field(default=600)
|
||
name: str = dataclasses.field(default="PiKVM")
|
||
|
||
|
||
class _Client(RfbClient): # pylint: disable=too-many-instance-attributes
|
||
def __init__( # pylint: disable=too-many-arguments,too-many-locals
|
||
self,
|
||
reader: asyncio.StreamReader,
|
||
writer: asyncio.StreamWriter,
|
||
tls_ciphers: str,
|
||
tls_timeout: float,
|
||
x509_cert_path: str,
|
||
x509_key_path: str,
|
||
|
||
desired_fps: int,
|
||
mouse_output: str,
|
||
keymap_name: str,
|
||
symmap: dict[int, dict[int, str]],
|
||
allow_cut_after: float,
|
||
|
||
kvmd: KvmdClient,
|
||
streamers: list[BaseStreamerClient],
|
||
|
||
vnc_credentials: dict[str, VncAuthKvmdCredentials],
|
||
vencrypt: bool,
|
||
none_auth_only: bool,
|
||
shared_params: _SharedParams,
|
||
) -> None:
|
||
|
||
self.__vnc_credentials = vnc_credentials
|
||
|
||
super().__init__(
|
||
reader=reader,
|
||
writer=writer,
|
||
tls_ciphers=tls_ciphers,
|
||
tls_timeout=tls_timeout,
|
||
x509_cert_path=x509_cert_path,
|
||
x509_key_path=x509_key_path,
|
||
allow_cut_after=allow_cut_after,
|
||
vnc_passwds=list(vnc_credentials),
|
||
vencrypt=vencrypt,
|
||
none_auth_only=none_auth_only,
|
||
**dataclasses.asdict(shared_params),
|
||
)
|
||
|
||
self.__desired_fps = desired_fps
|
||
self.__mouse_output = mouse_output
|
||
self.__keymap_name = keymap_name
|
||
self.__symmap = symmap
|
||
|
||
self.__kvmd = kvmd
|
||
self.__streamers = streamers
|
||
|
||
self.__shared_params = shared_params
|
||
|
||
self.__stage1_authorized = aiotools.AioStage()
|
||
self.__stage2_encodings_accepted = aiotools.AioStage()
|
||
self.__stage3_ws_connected = aiotools.AioStage()
|
||
|
||
self.__kvmd_session: (KvmdClientSession | None) = None
|
||
self.__kvmd_ws: (KvmdClientWs | None) = None
|
||
|
||
self.__fb_queue: "asyncio.Queue[dict]" = asyncio.Queue()
|
||
self.__fb_has_key = False
|
||
|
||
# Эти состояния шарить не обязательно - бекенд исключает дублирующиеся события.
|
||
# Все это нужно только чтобы не посылать лишние жсоны в сокет KVMD
|
||
self.__mouse_buttons: dict[str, (bool | None)] = dict.fromkeys(["left", "right", "middle", "up", "down"], None)
|
||
self.__mouse_move = {"x": -1, "y": -1}
|
||
|
||
self.__modifiers = 0
|
||
|
||
# =====
|
||
|
||
async def run(self) -> None:
|
||
try:
|
||
await self._run(
|
||
kvmd=self.__kvmd_task_loop(),
|
||
streamer=self.__streamer_task_loop(),
|
||
fb_sender=self.__fb_sender_task_loop(),
|
||
)
|
||
finally:
|
||
await aiotools.shield_fg(self.__cleanup())
|
||
|
||
async def __cleanup(self) -> None:
|
||
if self.__kvmd_session:
|
||
await self.__kvmd_session.close()
|
||
self.__kvmd_session = None
|
||
|
||
# =====
|
||
|
||
async def __kvmd_task_loop(self) -> None:
|
||
logger = get_logger(0)
|
||
await self.__stage1_authorized.wait_passed()
|
||
|
||
logger.info("%s [kvmd]: Waiting for the SetEncodings message ...", self._remote)
|
||
if not (await self.__stage2_encodings_accepted.wait_passed(timeout=5)):
|
||
raise RfbError("No SetEncodings message recieved from the client in 5 secs")
|
||
|
||
assert self.__kvmd_session
|
||
try:
|
||
logger.info("%s [kvmd]: Applying HID params: mouse_output=%s ...", self._remote, self.__mouse_output)
|
||
await self.__kvmd_session.hid.set_params(mouse_output=self.__mouse_output)
|
||
|
||
async with self.__kvmd_session.ws() as self.__kvmd_ws:
|
||
logger.info("%s [kvmd]: Connected to KVMD websocket", self._remote)
|
||
self.__stage3_ws_connected.set_passed()
|
||
async for (event_type, event) in self.__kvmd_ws.communicate():
|
||
await self.__process_ws_event(event_type, event)
|
||
raise RfbError("KVMD closed the websocket (the server may have been stopped)")
|
||
finally:
|
||
self.__kvmd_ws = None
|
||
|
||
async def __process_ws_event(self, event_type: str, event: dict) -> None:
|
||
if event_type == "info":
|
||
if "meta" in event:
|
||
try:
|
||
host = event["meta"]["server"]["host"]
|
||
except Exception:
|
||
host = None
|
||
else:
|
||
if isinstance(host, str):
|
||
name = f"PiKVM: {host}"
|
||
if self._encodings.has_rename:
|
||
await self._send_rename(name)
|
||
self.__shared_params.name = name
|
||
|
||
elif event_type == "hid":
|
||
if (
|
||
self._encodings.has_leds_state
|
||
and ("keyboard" in event)
|
||
and ("leds" in event["keyboard"])
|
||
):
|
||
await self._send_leds_state(**event["keyboard"]["leds"])
|
||
|
||
# =====
|
||
|
||
async def __streamer_task_loop(self) -> None:
|
||
logger = get_logger(0)
|
||
await self.__stage3_ws_connected.wait_passed()
|
||
streamer = self.__get_preferred_streamer()
|
||
while True:
|
||
try:
|
||
streaming = False
|
||
async with streamer.reading() as read_frame:
|
||
while True:
|
||
frame = await read_frame(not self.__fb_has_key)
|
||
if not streaming:
|
||
logger.info("%s [streamer]: Streaming ...", self._remote)
|
||
streaming = True
|
||
if frame["online"]:
|
||
await self.__queue_frame(frame)
|
||
else:
|
||
await self.__queue_frame("No signal")
|
||
except StreamerError as ex:
|
||
if isinstance(ex, StreamerPermError):
|
||
streamer = self.__get_default_streamer()
|
||
logger.info("%s [streamer]: Permanent error: %s; switching to %s ...", self._remote, ex, streamer)
|
||
else:
|
||
logger.info("%s [streamer]: Waiting for stream: %s", self._remote, ex)
|
||
await self.__queue_frame("Waiting for stream ...")
|
||
await asyncio.sleep(1)
|
||
|
||
def __get_preferred_streamer(self) -> BaseStreamerClient:
|
||
formats = {
|
||
StreamerFormats.JPEG: "has_tight",
|
||
StreamerFormats.H264: "has_h264",
|
||
}
|
||
streamer: (BaseStreamerClient | None) = None
|
||
for streamer in self.__streamers:
|
||
if getattr(self._encodings, formats[streamer.get_format()]):
|
||
get_logger(0).info("%s [streamer]: Using preferred %s", self._remote, streamer)
|
||
return streamer
|
||
raise RuntimeError("No streamers found")
|
||
|
||
def __get_default_streamer(self) -> BaseStreamerClient:
|
||
streamer = self.__streamers[-1]
|
||
get_logger(0).info("%s [streamer]: Using default %s", self._remote, streamer)
|
||
return streamer
|
||
|
||
async def __queue_frame(self, frame: (dict | str)) -> None:
|
||
if isinstance(frame, str):
|
||
frame = await self.__make_text_frame(frame)
|
||
if self.__fb_queue.qsize() > 10:
|
||
self.__fb_queue.get_nowait()
|
||
self.__fb_queue.put_nowait(frame)
|
||
|
||
async def __make_text_frame(self, text: str) -> dict:
|
||
return {
|
||
"data": (await make_text_jpeg(self._width, self._height, self._encodings.tight_jpeg_quality, text)),
|
||
"width": self._width,
|
||
"height": self._height,
|
||
"format": StreamerFormats.JPEG,
|
||
}
|
||
|
||
async def __fb_sender_task_loop(self) -> None: # pylint: disable=too-many-branches
|
||
last: (dict | None) = None
|
||
async for _ in self._send_fb_allowed():
|
||
while True:
|
||
frame = await self.__fb_queue.get()
|
||
if (
|
||
last is None # pylint: disable=too-many-boolean-expressions
|
||
or frame["format"] == StreamerFormats.JPEG
|
||
or last["format"] != frame["format"]
|
||
or (frame["format"] == StreamerFormats.H264 and (
|
||
frame["key"]
|
||
or last["width"] != frame["width"]
|
||
or last["height"] != frame["height"]
|
||
or len(last["data"]) + len(frame["data"]) > 4194304
|
||
))
|
||
):
|
||
self.__fb_has_key = (frame["format"] == StreamerFormats.H264 and frame["key"])
|
||
last = frame
|
||
if self.__fb_queue.qsize() == 0:
|
||
break
|
||
continue
|
||
assert frame["format"] == StreamerFormats.H264
|
||
last["data"] += frame["data"]
|
||
if self.__fb_queue.qsize() == 0:
|
||
break
|
||
|
||
if self._width != last["width"] or self._height != last["height"]:
|
||
self.__shared_params.width = last["width"]
|
||
self.__shared_params.height = last["height"]
|
||
if not self._encodings.has_resize:
|
||
msg = (
|
||
f"Resoultion changed: {self._width}x{self._height}"
|
||
f" -> {last['width']}x{last['height']}\nPlease reconnect"
|
||
)
|
||
await self._send_fb_jpeg((await self.__make_text_frame(msg))["data"])
|
||
continue
|
||
await self._send_resize(last["width"], last["height"])
|
||
|
||
if len(last["data"]) == 0:
|
||
# Вдруг какой-то баг
|
||
await self._send_fb_allow_again()
|
||
continue
|
||
|
||
if last["format"] == StreamerFormats.JPEG:
|
||
await self._send_fb_jpeg(last["data"])
|
||
elif last["format"] == StreamerFormats.H264:
|
||
if not self._encodings.has_h264:
|
||
raise RfbError("The client doesn't want to accept H264 anymore")
|
||
if self.__fb_has_key:
|
||
await self._send_fb_h264(last["data"])
|
||
else:
|
||
await self._send_fb_allow_again()
|
||
else:
|
||
raise RuntimeError(f"Unknown format: {last['format']}")
|
||
last["data"] = b""
|
||
|
||
# =====
|
||
|
||
async def _authorize_userpass(self, user: str, passwd: str) -> bool:
|
||
self.__kvmd_session = self.__kvmd.make_session(user, passwd)
|
||
if (await self.__kvmd_session.auth.check()):
|
||
self.__stage1_authorized.set_passed()
|
||
return True
|
||
return False
|
||
|
||
async def _on_authorized_vnc_passwd(self, passwd: str) -> str:
|
||
kc = self.__vnc_credentials[passwd]
|
||
if (await self._authorize_userpass(kc.user, kc.passwd)):
|
||
return kc.user
|
||
return ""
|
||
|
||
async def _on_authorized_none(self) -> bool:
|
||
return (await self._authorize_userpass("", ""))
|
||
|
||
# =====
|
||
|
||
async def _on_key_event(self, code: int, state: bool) -> None:
|
||
is_modifier = self.__switch_modifiers(code, state)
|
||
variants = self.__symmap.get(code)
|
||
fake_shift = False
|
||
|
||
if variants:
|
||
if is_modifier:
|
||
web_key = variants.get(0)
|
||
else:
|
||
web_key = variants.get(self.__modifiers)
|
||
if web_key is None:
|
||
web_key = variants.get(0)
|
||
|
||
if web_key is None and self.__modifiers == 0 and SymmapModifiers.SHIFT in variants:
|
||
# JUMP doesn't send shift events:
|
||
# - https://github.com/pikvm/pikvm/issues/820
|
||
web_key = variants[SymmapModifiers.SHIFT]
|
||
fake_shift = True
|
||
|
||
if web_key and self.__kvmd_ws:
|
||
if fake_shift:
|
||
await self.__kvmd_ws.send_key_event(WebModifiers.SHIFT_LEFT, True)
|
||
await self.__kvmd_ws.send_key_event(web_key, state)
|
||
if fake_shift:
|
||
await self.__kvmd_ws.send_key_event(WebModifiers.SHIFT_LEFT, False)
|
||
|
||
async def _on_ext_key_event(self, code: int, state: bool) -> None:
|
||
web_key = AT1_TO_WEB.get(code)
|
||
if web_key:
|
||
self.__switch_modifiers(web_key, state) # Предполагаем, что модификаторы всегда известны
|
||
if self.__kvmd_ws:
|
||
await self.__kvmd_ws.send_key_event(web_key, state)
|
||
|
||
def __switch_modifiers(self, key: (int | str), state: bool) -> bool:
|
||
mod = 0
|
||
if key in X11Modifiers.SHIFTS or key in WebModifiers.SHIFTS:
|
||
mod = SymmapModifiers.SHIFT
|
||
elif key == X11Modifiers.ALTGR or key == WebModifiers.ALT_RIGHT:
|
||
mod = SymmapModifiers.ALTGR
|
||
elif key in X11Modifiers.CTRLS or key in WebModifiers.CTRLS:
|
||
mod = SymmapModifiers.CTRL
|
||
if mod == 0:
|
||
return False
|
||
if state:
|
||
self.__modifiers |= mod
|
||
else:
|
||
self.__modifiers &= ~mod
|
||
return True
|
||
|
||
async def _on_pointer_event(self, buttons: dict[str, bool], wheel: dict[str, int], move: dict[str, int]) -> None:
|
||
if self.__kvmd_ws:
|
||
if wheel["x"] or wheel["y"]:
|
||
await self.__kvmd_ws.send_mouse_wheel_event(wheel["x"], wheel["y"])
|
||
|
||
if self.__mouse_move != move:
|
||
await self.__kvmd_ws.send_mouse_move_event(move["x"], move["y"])
|
||
self.__mouse_move = move
|
||
|
||
for (button, state) in buttons.items():
|
||
if self.__mouse_buttons[button] != state:
|
||
await self.__kvmd_ws.send_mouse_button_event(button, state)
|
||
self.__mouse_buttons[button] = state
|
||
|
||
async def _on_cut_event(self, text: str) -> None:
|
||
assert self.__stage1_authorized.is_passed()
|
||
assert self.__kvmd_session
|
||
logger = get_logger(0)
|
||
logger.info("%s [main]: Printing %d characters ...", self._remote, len(text))
|
||
try:
|
||
(keymap_name, available) = await self.__kvmd_session.hid.get_keymaps()
|
||
if self.__keymap_name in available:
|
||
keymap_name = self.__keymap_name
|
||
await self.__kvmd_session.hid.print(text, 0, keymap_name)
|
||
except Exception:
|
||
logger.exception("%s [main]: Can't print characters", self._remote)
|
||
|
||
async def _on_set_encodings(self) -> None:
|
||
assert self.__stage1_authorized.is_passed()
|
||
assert self.__kvmd_session
|
||
self.__stage2_encodings_accepted.set_passed(multi=True)
|
||
|
||
has_quality = (await self.__kvmd_session.streamer.get_state())["features"]["quality"]
|
||
quality = (self._encodings.tight_jpeg_quality if has_quality else None)
|
||
get_logger(0).info("%s [main]: Applying streamer params: jpeg_quality=%s; desired_fps=%d ...",
|
||
self._remote, quality, self.__desired_fps)
|
||
await self.__kvmd_session.streamer.set_params(quality, self.__desired_fps)
|
||
|
||
|
||
# =====
|
||
class VncServer: # pylint: disable=too-many-instance-attributes
|
||
def __init__( # pylint: disable=too-many-arguments,too-many-locals
|
||
self,
|
||
host: str,
|
||
port: int,
|
||
max_clients: int,
|
||
|
||
no_delay: bool,
|
||
keepalive_enabled: bool,
|
||
keepalive_idle: int,
|
||
keepalive_interval: int,
|
||
keepalive_count: int,
|
||
|
||
tls_ciphers: str,
|
||
tls_timeout: float,
|
||
x509_cert_path: str,
|
||
x509_key_path: str,
|
||
|
||
vencrypt_enabled: bool,
|
||
|
||
desired_fps: int,
|
||
mouse_output: str,
|
||
keymap_path: str,
|
||
allow_cut_after: float,
|
||
|
||
kvmd: KvmdClient,
|
||
streamers: list[BaseStreamerClient],
|
||
vnc_auth_manager: VncAuthManager,
|
||
) -> None:
|
||
|
||
self.__host = network.get_listen_host(host)
|
||
self.__port = port
|
||
self.__max_clients = max_clients
|
||
|
||
keymap_name = os.path.basename(keymap_path)
|
||
symmap = build_symmap(keymap_path)
|
||
|
||
self.__vnc_auth_manager = vnc_auth_manager
|
||
|
||
shared_params = _SharedParams()
|
||
|
||
async def cleanup_client(writer: asyncio.StreamWriter) -> None:
|
||
if (await aiotools.close_writer(writer)):
|
||
get_logger(0).info("%s [entry]: Connection is closed in an emergency", rfb_format_remote(writer))
|
||
|
||
async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
||
logger = get_logger(0)
|
||
remote = rfb_format_remote(writer)
|
||
logger.info("%s [entry]: Connected client", remote)
|
||
try:
|
||
sock = writer.get_extra_info("socket")
|
||
if no_delay:
|
||
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
||
if keepalive_enabled:
|
||
# https://www.tldp.org/HOWTO/html_single/TCP-Keepalive-HOWTO/#setsockopt
|
||
# https://blog.cloudflare.com/when-tcp-sockets-refuse-to-die
|
||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
|
||
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, keepalive_idle)
|
||
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, keepalive_interval)
|
||
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, keepalive_count)
|
||
timeout = (keepalive_idle + keepalive_interval * keepalive_count) * 1000 # Milliseconds
|
||
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT, timeout) # type: ignore
|
||
|
||
try:
|
||
async with kvmd.make_session("", "") as kvmd_session:
|
||
none_auth_only = await kvmd_session.auth.check()
|
||
except (aiohttp.ClientError, asyncio.TimeoutError) as ex:
|
||
logger.error("%s [entry]: Can't check KVMD auth mode: %s", remote, tools.efmt(ex))
|
||
return
|
||
|
||
await _Client(
|
||
reader=reader,
|
||
writer=writer,
|
||
tls_ciphers=tls_ciphers,
|
||
tls_timeout=tls_timeout,
|
||
x509_cert_path=x509_cert_path,
|
||
x509_key_path=x509_key_path,
|
||
desired_fps=desired_fps,
|
||
mouse_output=mouse_output,
|
||
keymap_name=keymap_name,
|
||
symmap=symmap,
|
||
allow_cut_after=allow_cut_after,
|
||
kvmd=kvmd,
|
||
streamers=streamers,
|
||
vnc_credentials=(await self.__vnc_auth_manager.read_credentials())[0],
|
||
none_auth_only=none_auth_only,
|
||
vencrypt=vencrypt_enabled,
|
||
shared_params=shared_params,
|
||
).run()
|
||
except Exception:
|
||
logger.exception("%s [entry]: Unhandled exception in client task", remote)
|
||
finally:
|
||
await aiotools.shield_fg(cleanup_client(writer))
|
||
|
||
self.__handle_client = handle_client
|
||
|
||
async def __inner_run(self) -> None:
|
||
if not (await self.__vnc_auth_manager.read_credentials())[1]:
|
||
raise SystemExit(1)
|
||
|
||
get_logger(0).info("Listening VNC on TCP [%s]:%d ...", self.__host, self.__port)
|
||
(family, _, _, _, addr) = socket.getaddrinfo(self.__host, self.__port, type=socket.SOCK_STREAM)[0]
|
||
with contextlib.closing(socket.socket(family, socket.SOCK_STREAM)) as sock:
|
||
if family == socket.AF_INET6:
|
||
sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
|
||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||
sock.bind(addr)
|
||
|
||
server = await asyncio.start_server(
|
||
client_connected_cb=self.__handle_client,
|
||
sock=sock,
|
||
backlog=self.__max_clients,
|
||
)
|
||
async with server:
|
||
await server.serve_forever()
|
||
|
||
def run(self) -> None:
|
||
aiotools.run(self.__inner_run())
|
||
get_logger().info("Bye-bye")
|