From 16ad64db886296c58adc3ec8b0f266e57f75fb1b Mon Sep 17 00:00:00 2001 From: Devaev Maxim Date: Sun, 22 Nov 2020 15:36:45 +0300 Subject: refactoring --- kvmd/plugins/msd/relay/__init__.py | 279 +++++++++++++++++++++++++++++++++++++ kvmd/plugins/msd/relay/drive.py | 143 +++++++++++++++++++ kvmd/plugins/msd/relay/gpio.py | 78 +++++++++++ 3 files changed, 500 insertions(+) create mode 100644 kvmd/plugins/msd/relay/__init__.py create mode 100644 kvmd/plugins/msd/relay/drive.py create mode 100644 kvmd/plugins/msd/relay/gpio.py (limited to 'kvmd/plugins/msd/relay') diff --git a/kvmd/plugins/msd/relay/__init__.py b/kvmd/plugins/msd/relay/__init__.py new file mode 100644 index 00000000..ea87168c --- /dev/null +++ b/kvmd/plugins/msd/relay/__init__.py @@ -0,0 +1,279 @@ +# ========================================================================== # +# # +# KVMD - The main Pi-KVM daemon. # +# # +# Copyright (C) 2018 Maxim Devaev # +# # +# 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 . # +# # +# ========================================================================== # + + +import asyncio +import contextlib +import dataclasses + +from typing import Dict +from typing import AsyncGenerator +from typing import Optional + +import aiofiles +import aiofiles.base + +from ....logging import get_logger + +from .... import aiotools +from .... import aiofs + +from ....yamlconf import Option + +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 BaseMsd + +from .gpio import Gpio + +from .drive import ImageInfo +from .drive import DeviceInfo + + +# ===== +class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes + def __init__( # pylint: disable=super-init-not-called + self, + target_pin: int, + reset_pin: int, + + device_path: str, + init_delay: float, + init_retries: int, + reset_delay: float, + ) -> None: + + self.__device_path = device_path + self.__init_delay = init_delay + self.__init_retries = init_retries + + self.__gpio = Gpio(target_pin, reset_pin, reset_delay) + + self.__device_info: Optional[DeviceInfo] = None + self.__connected = False + + self.__device_file: Optional[aiofiles.base.AiofilesContextManager] = None + self.__written = 0 + + self.__notifier = aiotools.AioNotifier() + self.__region = aiotools.AioExclusiveRegion(MsdIsBusyError, self.__notifier) + + logger = get_logger(0) + 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) + + @classmethod + def get_plugin_options(cls) -> Dict: + return { + "target_pin": Option(-1, type=valid_gpio_pin), + "reset_pin": Option(-1, type=valid_gpio_pin), + + "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: + self.__gpio.open() + + 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": bool(self.__device_file) + } + 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, + }, + } + + 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_file() + finally: + self.__gpio.close() + + # ===== + + @aiotools.atomic + async def set_params(self, name: Optional[str]=None, cdrom: Optional[bool]=None) -> None: + async with self.__working(): + if name is not None: + raise MsdMultiNotSupported() + if cdrom is not None: + raise MsdCdromNotSupported() + + @aiotools.atomic + async def connect(self) -> None: + async with self.__working(): + async with self.__region: + if self.__connected: + raise MsdConnectedError() + + self.__gpio.switch_to_server() + self.__connected = True + get_logger(0).info("MSD switched to Server") + + @aiotools.atomic + async def disconnect(self) -> None: + async with self.__working(): + async with self.__region: + 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 + self.__connected = False + get_logger(0).info("MSD switched to KVM: %s", self.__device_info) + + @contextlib.asynccontextmanager + async def write_image(self, name: str) -> AsyncGenerator[None, None]: + async with self.__working(): + async with self.__region: + try: + assert self.__device_info + if self.__connected: + raise MsdConnectedError() + + self.__device_file = await aiofiles.open(self.__device_info.path, mode="w+b", buffering=0) # type: ignore + self.__written = 0 + + await self.__write_image_info(name, complete=False) + await self.__notifier.notify() + yield + await self.__write_image_info(name, complete=True) + finally: + await self.__close_device_file() + await self.__load_device_info() + + async def write_image_chunk(self, chunk: bytes) -> int: + assert self.__device_file + await aiofs.afile_write_now(self.__device_file, chunk) + self.__written += len(chunk) + return self.__written + + @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, name: str, complete: bool) -> None: + assert self.__device_file + assert self.__device_info + if not self.__device_info.write_image_info( + device_file=self.__device_file, + image_info=ImageInfo(name, self.__written, complete), + ): + get_logger().error("Can't write image info because device is full") + + async def __close_device_file(self) -> None: + try: + if self.__device_file: + get_logger().info("Closing device file ...") + await self.__device_file.close() # type: ignore + except Exception: + get_logger().exception("Can't close device file") + finally: + self.__device_file = None + self.__written = 0 + + 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 diff --git a/kvmd/plugins/msd/relay/drive.py b/kvmd/plugins/msd/relay/drive.py new file mode 100644 index 00000000..a04a6e9f --- /dev/null +++ b/kvmd/plugins/msd/relay/drive.py @@ -0,0 +1,143 @@ +# ========================================================================== # +# # +# KVMD - The main Pi-KVM daemon. # +# # +# Copyright (C) 2018 Maxim Devaev # +# # +# 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 . # +# # +# ========================================================================== # + + +import os +import stat +import fcntl +import struct +import dataclasses + +from typing import IO +from typing import Optional + +import aiofiles.base + +from .... import aiotools +from .... import aiofs + + +# ===== +_IMAGE_INFO_SIZE = 4096 +_IMAGE_INFO_MAGIC_SIZE = 16 +_IMAGE_INFO_NAME_SIZE = 256 +_IMAGE_INFO_PADS_SIZE = _IMAGE_INFO_SIZE - _IMAGE_INFO_NAME_SIZE - 1 - 8 - _IMAGE_INFO_MAGIC_SIZE * 8 +_IMAGE_INFO_FORMAT = ">%dL%dc?Q%dx%dL" % ( + _IMAGE_INFO_MAGIC_SIZE, + _IMAGE_INFO_NAME_SIZE, + _IMAGE_INFO_PADS_SIZE, + _IMAGE_INFO_MAGIC_SIZE, +) +_IMAGE_INFO_MAGIC = [0x1ACE1ACE] * _IMAGE_INFO_MAGIC_SIZE + + +# ===== +@dataclasses.dataclass(frozen=True) +class ImageInfo: + name: str + size: int + complete: bool + + @classmethod + def from_bytes(cls, data: bytes) -> Optional["ImageInfo"]: + try: + parsed = list(struct.unpack(_IMAGE_INFO_FORMAT, data)) + except struct.error: + pass + else: + magic_begin = parsed[:_IMAGE_INFO_MAGIC_SIZE] + magic_end = parsed[-_IMAGE_INFO_MAGIC_SIZE:] + if magic_begin == magic_end == _IMAGE_INFO_MAGIC: + image_name_bytes = b"".join(parsed[ + _IMAGE_INFO_MAGIC_SIZE # noqa: E203 + : + _IMAGE_INFO_MAGIC_SIZE + _IMAGE_INFO_NAME_SIZE + ]) + return ImageInfo( + name=image_name_bytes.decode("utf-8", errors="ignore").strip("\x00").strip(), + size=parsed[_IMAGE_INFO_MAGIC_SIZE + _IMAGE_INFO_NAME_SIZE + 1], + complete=parsed[_IMAGE_INFO_MAGIC_SIZE + _IMAGE_INFO_NAME_SIZE], + ) + return None + + def to_bytes(self) -> bytes: + return struct.pack( + _IMAGE_INFO_FORMAT, + *_IMAGE_INFO_MAGIC, + *memoryview(( # type: ignore + self.name.encode("utf-8") + + b"\x00" * _IMAGE_INFO_NAME_SIZE + )[:_IMAGE_INFO_NAME_SIZE]).cast("c"), + self.complete, + self.size, + *_IMAGE_INFO_MAGIC, + ) + + +@dataclasses.dataclass(frozen=True) +class DeviceInfo: + path: str + size: int + free: int + image: Optional[ImageInfo] + + @classmethod + async def read(cls, device_path: str) -> "DeviceInfo": + return (await aiotools.run_async(cls.__inner_read, device_path)) + + @classmethod + def __inner_read(cls, device_path: str) -> "DeviceInfo": + if not stat.S_ISBLK(os.stat(device_path).st_mode): + raise RuntimeError(f"Not a block device: {device_path}") + + with open(device_path, "rb") as device_file: + # size = BLKGETSIZE * BLKSSZGET + size = _ioctl_uint32(device_file, 0x1260) * _ioctl_uint32(device_file, 0x1268) + device_file.seek(size - _IMAGE_INFO_SIZE) + image_info = ImageInfo.from_bytes(device_file.read()) + + return DeviceInfo( + path=device_path, + size=size, + free=(size - image_info.size if image_info else size), + image=image_info, + ) + + async def write_image_info( + self, + device_file: aiofiles.base.AiofilesContextManager, + image_info: ImageInfo, + ) -> bool: + + if self.size - image_info.size > _IMAGE_INFO_SIZE: + await device_file.seek(self.size - _IMAGE_INFO_SIZE) # type: ignore + await aiofs.afile_write_now(device_file, image_info.to_bytes()) + await device_file.seek(0) # type: ignore + return True + return False # Device is full + + +def _ioctl_uint32(device_file: IO, request: int) -> int: + buf = b"\0" * 4 + buf = fcntl.ioctl(device_file.fileno(), request, buf) # type: ignore + result = struct.unpack("I", buf)[0] + assert result > 0, (device_file, request, buf) + return result diff --git a/kvmd/plugins/msd/relay/gpio.py b/kvmd/plugins/msd/relay/gpio.py new file mode 100644 index 00000000..0950c1d6 --- /dev/null +++ b/kvmd/plugins/msd/relay/gpio.py @@ -0,0 +1,78 @@ +# ========================================================================== # +# # +# KVMD - The main Pi-KVM daemon. # +# # +# Copyright (C) 2018 Maxim Devaev # +# # +# 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 . # +# # +# ========================================================================== # + + +from typing import Optional + +import gpiod + +from .... import env +from .... import aiogp + + +# ===== +class Gpio: + def __init__( + self, + target_pin: int, + reset_pin: int, + reset_delay: float, + ) -> None: + + self.__target_pin = target_pin + self.__reset_pin = reset_pin + self.__reset_delay = reset_delay + + self.__chip: Optional[gpiod.Chip] = None + self.__target_line: Optional[gpiod.Line] = None + self.__reset_line: Optional[gpiod.Line] = None + + def open(self) -> None: + assert self.__chip is None + assert self.__target_line is None + assert self.__reset_line is None + + self.__chip = gpiod.Chip(env.GPIO_DEVICE_PATH) + + self.__target_line = self.__chip.get_line(self.__target_pin) + self.__target_line.request("kvmd::msd::target", gpiod.LINE_REQ_DIR_OUT, default_vals=[0]) + + self.__reset_line = self.__chip.get_line(self.__reset_pin) + self.__reset_line.request("kvmd::msd::reset", gpiod.LINE_REQ_DIR_OUT, default_vals=[0]) + + def close(self) -> None: + if self.__chip: + try: + self.__chip.close() + except Exception: + pass + + def switch_to_local(self) -> None: + assert self.__target_line + self.__target_line.set_value(0) + + def switch_to_server(self) -> None: + assert self.__target_line + self.__target_line.set_value(1) + + async def reset(self) -> None: + assert self.__reset_line + await aiogp.pulse(self.__reset_line, self.__reset_delay, 0) -- cgit v1.2.3