mirror of
https://github.com/mofeng-git/One-KVM.git
synced 2025-12-12 17:20:30 +08:00
300 lines
11 KiB
Python
300 lines
11 KiB
Python
# ========================================================================== #
|
|
# #
|
|
# KVMD - The main PiKVM daemon. #
|
|
# #
|
|
# Copyright (C) 2018-2022 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 asyncio
|
|
import contextlib
|
|
import dataclasses
|
|
import functools
|
|
|
|
from typing import Dict
|
|
from typing import AsyncGenerator
|
|
from typing import Optional
|
|
|
|
from ....logging import get_logger
|
|
|
|
from .... import aiotools
|
|
|
|
from ....yamlconf import Option
|
|
|
|
from ....validators.basic import valid_bool
|
|
from ....validators.basic import valid_number
|
|
from ....validators.basic import valid_int_f1
|
|
from ....validators.basic import valid_float_f01
|
|
from ....validators.os import valid_abs_path
|
|
from ....validators.hw import valid_gpio_pin
|
|
|
|
from .. import MsdError
|
|
from .. import MsdIsBusyError
|
|
from .. import MsdOfflineError
|
|
from .. import MsdConnectedError
|
|
from .. import MsdDisconnectedError
|
|
from .. import MsdMultiNotSupported
|
|
from .. import MsdCdromNotSupported
|
|
from .. import MsdRwNotSupported
|
|
from .. import BaseMsd
|
|
from .. import MsdImageWriter
|
|
|
|
from .gpio import Gpio
|
|
|
|
from .drive import DeviceInfo
|
|
|
|
|
|
# =====
|
|
class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
|
|
def __init__( # pylint: disable=super-init-not-called,too-many-arguments
|
|
self,
|
|
upload_chunk_size: int,
|
|
sync_chunk_size: int,
|
|
|
|
gpio_device_path: str,
|
|
target_pin: int,
|
|
reset_inverted: bool,
|
|
reset_pin: int,
|
|
|
|
device_path: str,
|
|
init_delay: float,
|
|
init_retries: int,
|
|
reset_delay: float,
|
|
) -> None:
|
|
|
|
self.__upload_chunk_size = upload_chunk_size
|
|
self.__sync_chunk_size = sync_chunk_size
|
|
|
|
self.__device_path = device_path
|
|
self.__init_delay = init_delay
|
|
self.__init_retries = init_retries
|
|
|
|
self.__gpio = Gpio(gpio_device_path, target_pin, reset_pin, reset_inverted, reset_delay)
|
|
|
|
self.__device_info: Optional[DeviceInfo] = None
|
|
self.__connected = False
|
|
|
|
self.__device_writer: Optional[MsdImageWriter] = None
|
|
|
|
self.__notifier = aiotools.AioNotifier()
|
|
self.__region = aiotools.AioExclusiveRegion(MsdIsBusyError, self.__notifier)
|
|
|
|
@classmethod
|
|
def get_plugin_options(cls) -> Dict:
|
|
return {
|
|
"upload_chunk_size": Option(65536, type=functools.partial(valid_number, min=1024)),
|
|
"sync_chunk_size": Option(4194304, type=functools.partial(valid_number, min=1024)),
|
|
|
|
"gpio_device": Option("/dev/gpiochip0", type=valid_abs_path, unpack_as="gpio_device_path"),
|
|
"target_pin": Option(-1, type=valid_gpio_pin),
|
|
"reset_pin": Option(-1, type=valid_gpio_pin),
|
|
"reset_inverted": Option(False, type=valid_bool),
|
|
|
|
"device": Option("", type=valid_abs_path, unpack_as="device_path"),
|
|
"init_delay": Option(1.0, type=valid_float_f01),
|
|
"init_retries": Option(5, type=valid_int_f1),
|
|
"reset_delay": Option(1.0, type=valid_float_f01),
|
|
}
|
|
|
|
def sysprep(self) -> None:
|
|
logger = get_logger(0)
|
|
self.__gpio.open()
|
|
logger.info("Using %r as MSD", self.__device_path)
|
|
try:
|
|
aiotools.run_sync(self.__load_device_info())
|
|
except Exception as err:
|
|
log = (logger.error if isinstance(err, MsdError) else logger.exception)
|
|
log("MSD is offline: %s", err)
|
|
|
|
async def get_state(self) -> Dict:
|
|
storage: Optional[Dict] = None
|
|
drive: Optional[Dict] = None
|
|
if self.__device_info:
|
|
storage = {
|
|
"size": self.__device_info.size,
|
|
"free": self.__device_info.free,
|
|
"uploading": (self.__device_writer.get_state() if self.__device_writer else None),
|
|
}
|
|
drive = {
|
|
"image": (self.__device_info.image and dataclasses.asdict(self.__device_info.image)),
|
|
"connected": self.__connected,
|
|
}
|
|
return {
|
|
"enabled": True,
|
|
"online": bool(self.__device_info),
|
|
"busy": self.__region.is_busy(),
|
|
"storage": storage,
|
|
"drive": drive,
|
|
"features": {
|
|
"multi": False,
|
|
"cdrom": False,
|
|
"rw": False,
|
|
},
|
|
}
|
|
|
|
async def poll_state(self) -> AsyncGenerator[Dict, None]:
|
|
prev_state: Dict = {}
|
|
while True:
|
|
state = await self.get_state()
|
|
if state != prev_state:
|
|
yield state
|
|
prev_state = state
|
|
await self.__notifier.wait()
|
|
|
|
@aiotools.atomic
|
|
async def reset(self) -> None:
|
|
await aiotools.run_region_task(
|
|
"Can't reset MSD or operation was not completed",
|
|
self.__region, self.__inner_reset,
|
|
)
|
|
|
|
@aiotools.atomic
|
|
async def __inner_reset(self) -> None:
|
|
await self.__gpio.reset()
|
|
self.__gpio.switch_to_local()
|
|
self.__connected = False
|
|
await self.__load_device_info()
|
|
get_logger(0).info("MSD reset has been successful")
|
|
|
|
@aiotools.atomic
|
|
async def cleanup(self) -> None:
|
|
try:
|
|
await self.__close_device_writer()
|
|
finally:
|
|
self.__gpio.close()
|
|
|
|
# =====
|
|
|
|
@aiotools.atomic
|
|
async def set_params(
|
|
self,
|
|
name: Optional[str]=None,
|
|
cdrom: Optional[bool]=None,
|
|
rw: Optional[bool]=None,
|
|
) -> None:
|
|
|
|
async with self.__working():
|
|
if name is not None:
|
|
raise MsdMultiNotSupported()
|
|
if cdrom is not None:
|
|
raise MsdCdromNotSupported()
|
|
if rw is not None:
|
|
raise MsdRwNotSupported()
|
|
|
|
@aiotools.atomic
|
|
async def set_connected(self, connected: bool) -> None:
|
|
async with self.__working():
|
|
async with self.__region:
|
|
if connected:
|
|
if self.__connected:
|
|
raise MsdConnectedError()
|
|
self.__gpio.switch_to_server()
|
|
get_logger(0).info("MSD switched to Server")
|
|
else:
|
|
if not self.__connected:
|
|
raise MsdDisconnectedError()
|
|
self.__gpio.switch_to_local()
|
|
try:
|
|
await self.__load_device_info()
|
|
except Exception:
|
|
if self.__connected:
|
|
self.__gpio.switch_to_server()
|
|
raise
|
|
get_logger(0).info("MSD switched to KVM: %s", self.__device_info)
|
|
self.__connected = connected
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def read_image(self, name: str) -> AsyncGenerator[int, None]:
|
|
async with self.__working():
|
|
if self is not None: # XXX: Vulture and pylint hack
|
|
raise MsdMultiNotSupported()
|
|
yield 1
|
|
|
|
async def read_image_chunk(self) -> bytes:
|
|
async with self.__working():
|
|
raise MsdMultiNotSupported()
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
|
|
async with self.__working():
|
|
if remove_incomplete is not None:
|
|
raise MsdMultiNotSupported()
|
|
async with self.__region:
|
|
try:
|
|
assert self.__device_info
|
|
if self.__connected:
|
|
raise MsdConnectedError()
|
|
|
|
self.__device_writer = await MsdImageWriter(
|
|
notifier=self.__notifier,
|
|
path=self.__device_info.path,
|
|
size=size,
|
|
sync=self.__sync_chunk_size,
|
|
).open()
|
|
|
|
await self.__write_image_info(False)
|
|
await self.__notifier.notify()
|
|
yield self.__upload_chunk_size
|
|
await self.__write_image_info(True)
|
|
finally:
|
|
await self.__close_device_writer()
|
|
await self.__load_device_info()
|
|
|
|
async def write_image_chunk(self, chunk: bytes) -> int:
|
|
assert self.__device_writer
|
|
return (await self.__device_writer.write(chunk))
|
|
|
|
@aiotools.atomic
|
|
async def remove(self, name: str) -> None:
|
|
async with self.__working():
|
|
raise MsdMultiNotSupported()
|
|
|
|
# =====
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def __working(self) -> AsyncGenerator[None, None]:
|
|
if not self.__device_info:
|
|
raise MsdOfflineError()
|
|
yield
|
|
|
|
# =====
|
|
|
|
async def __write_image_info(self, complete: bool) -> None:
|
|
assert self.__device_writer
|
|
assert self.__device_info
|
|
if not (await self.__device_info.write_image_info(self.__device_writer, complete)):
|
|
get_logger().error("Can't write image info because device is full")
|
|
|
|
async def __close_device_writer(self) -> None:
|
|
if self.__device_writer:
|
|
await self.__device_writer.close() # type: ignore
|
|
self.__device_writer = None
|
|
|
|
async def __load_device_info(self) -> None:
|
|
retries = self.__init_retries
|
|
while True:
|
|
await asyncio.sleep(self.__init_delay)
|
|
try:
|
|
self.__device_info = await DeviceInfo.read(self.__device_path)
|
|
break
|
|
except Exception:
|
|
if retries == 0:
|
|
self.__device_info = None
|
|
raise MsdError("Can't load device info")
|
|
get_logger().exception("Can't load device info; retries=%d", retries)
|
|
retries -= 1
|