mirror of
https://github.com/mofeng-git/One-KVM.git
synced 2026-02-01 02:21:53 +08:00
refactoring
This commit is contained in:
Devaev Maxim
279
kvmd/plugins/msd/relay/__init__.py
Normal file
279
kvmd/plugins/msd/relay/__init__.py
Normal file
@@ -0,0 +1,279 @@
|
||||
# ========================================================================== #
|
||||
# #
|
||||
# KVMD - The main Pi-KVM daemon. #
|
||||
# #
|
||||
# Copyright (C) 2018 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
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user