GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2025-12-12 17:20:30 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactoring

Browse Source
This commit is contained in:
Devaev Maxim 2018-11-28 23:41:03 +03:00
parent 5407f983c8
commit 5f7834724a
2 changed files with 76 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

72
kvmd/apps/kvmd/msd.py
View File

@ -1,12 +1,14 @@
import os import os
import struct import struct
import asyncio import asyncio
import asyncio.queues
import types import types
from typing import Dict from typing import Dict
from typing import NamedTuple from typing import NamedTuple
from typing import Callable from typing import Callable
from typing import Type from typing import Type
from typing import AsyncGenerator
from typing import Optional from typing import Optional
from typing import Any from typing import Any
@ -192,6 +194,8 @@ class MassStorageDevice: # pylint: disable=too-many-instance-attributes
self.__device_file: Optional[aiofiles.base.AiofilesContextManager] = None self.__device_file: Optional[aiofiles.base.AiofilesContextManager] = None
self.__written = 0 self.__written = 0
self.__state_queue: asyncio.queues.Queue = asyncio.Queue()
logger = get_logger(0) logger = get_logger(0)
if self._device_path: if self._device_path:
logger.info("Using %r as mass-storage device", self._device_path) logger.info("Using %r as mass-storage device", self._device_path)
@ -208,33 +212,6 @@ class MassStorageDevice: # pylint: disable=too-many-instance-attributes
else: else:
logger.warning("Mass-storage device is not operational") logger.warning("Mass-storage device is not operational")
@_msd_operated
async def connect_to_kvm(self, no_delay: bool=False) -> None:
with self.__region:
if self.__device_info:
raise MsdAlreadyConnectedToKvmError()
gpio.write(self.__target, False)
if not no_delay:
await asyncio.sleep(self.__init_delay)
await self.__load_device_info()
get_logger().info("Mass-storage device switched to KVM: %s", self.__device_info)
@_msd_operated
async def connect_to_pc(self) -> None:
with self.__region:
if not self.__device_info:
raise MsdAlreadyConnectedToPcError()
gpio.write(self.__target, True)
self.__device_info = None
get_logger().info("Mass-storage device switched to Server")
@_msd_operated
async def reset(self) -> None:
with self.__region:
gpio.write(self.__reset, True)
await asyncio.sleep(self.__reset_delay)
gpio.write(self.__reset, False)
def get_state(self) -> Dict: def get_state(self) -> Dict:
info = (self.__saved_device_info._asdict() if self.__saved_device_info else None) info = (self.__saved_device_info._asdict() if self.__saved_device_info else None)
if info: if info:
@ -253,11 +230,50 @@ class MassStorageDevice: # pylint: disable=too-many-instance-attributes
"info": info, "info": info,
} }
async def poll_state(self) -> AsyncGenerator[Dict, None]:
while True:
yield (await self.__state_queue.get())
async def cleanup(self) -> None: async def cleanup(self) -> None:
await self.__close_device_file() await self.__close_device_file()
gpio.write(self.__target, False) gpio.write(self.__target, False)
gpio.write(self.__reset, False) gpio.write(self.__reset, False)
@_msd_operated
async def connect_to_kvm(self, no_delay: bool=False) -> Dict:
with self.__region:
if self.__device_info:
raise MsdAlreadyConnectedToKvmError()
gpio.write(self.__target, False)
if not no_delay:
await asyncio.sleep(self.__init_delay)
await self.__load_device_info()
state = self.get_state()
await self.__state_queue.put(state)
get_logger().info("Mass-storage device switched to KVM: %s", self.__device_info)
return state
@_msd_operated
async def connect_to_pc(self) -> Dict:
with self.__region:
if not self.__device_info:
raise MsdAlreadyConnectedToPcError()
gpio.write(self.__target, True)
self.__device_info = None
state = self.get_state()
await self.__state_queue.put(state)
get_logger().info("Mass-storage device switched to Server")
return state
@_msd_operated
async def reset(self) -> None:
with self.__region:
get_logger().info("Mass-storage device reset")
gpio.write(self.__reset, True)
await asyncio.sleep(self.__reset_delay)
gpio.write(self.__reset, False)
await self.__state_queue.put(self.get_state())
@_msd_operated @_msd_operated
async def __aenter__(self) -> "MassStorageDevice": async def __aenter__(self) -> "MassStorageDevice":
self.__region.enter() self.__region.enter()
@ -268,6 +284,7 @@ class MassStorageDevice: # pylint: disable=too-many-instance-attributes
self.__written = 0 self.__written = 0
return self return self
finally: finally:
await self.__state_queue.put(self.get_state())
self.__region.exit() self.__region.exit()
async def write_image_info(self, name: str, complete: bool) -> None: async def write_image_info(self, name: str, complete: bool) -> None:
@ -296,6 +313,7 @@ class MassStorageDevice: # pylint: disable=too-many-instance-attributes
try: try:
await self.__close_device_file() await self.__close_device_file()
finally: finally:
await self.__state_queue.put(self.get_state())
self.__region.exit() self.__region.exit()
async def __write_to_device_file(self, data: bytes) -> None: async def __write_to_device_file(self, data: bytes) -> None:

66
kvmd/apps/kvmd/server.py
View File

@ -183,6 +183,7 @@ class Server: # pylint: disable=too-many-instance-attributes
self.__loop.create_task(self.__stream_controller()), self.__loop.create_task(self.__stream_controller()),
self.__loop.create_task(self.__poll_dead_sockets()), self.__loop.create_task(self.__poll_dead_sockets()),
self.__loop.create_task(self.__poll_atx_state()), self.__loop.create_task(self.__poll_atx_state()),
self.__loop.create_task(self.__poll_msd_state()),
self.__loop.create_task(self.__poll_streamer_state()), self.__loop.create_task(self.__poll_streamer_state()),
]) ])
@ -303,9 +304,7 @@ class Server: # pylint: disable=too-many-instance-attributes
}.get(button) }.get(button)
if not clicker: if not clicker:
raise BadRequest("Invalid param 'button'") raise BadRequest("Invalid param 'button'")
await self.__broadcast_event(_Events.ATX_STATE, self.__atx.get_state())
await clicker() await clicker()
await self.__broadcast_event(_Events.ATX_STATE, self.__atx.get_state())
return _json({"clicked": button}) return _json({"clicked": button})
# ===== MSD # ===== MSD
@ -317,16 +316,11 @@ class Server: # pylint: disable=too-many-instance-attributes
async def __msd_connect_handler(self, request: aiohttp.web.Request) -> aiohttp.web.Response: async def __msd_connect_handler(self, request: aiohttp.web.Request) -> aiohttp.web.Response:
to = request.query.get("to") to = request.query.get("to")
if to == "kvm": if to == "kvm":
await self.__msd.connect_to_kvm() return _json(await self.__msd.connect_to_kvm())
state = self.__msd.get_state()
await self.__broadcast_event(_Events.MSD_STATE, state)
elif to == "server": elif to == "server":
await self.__msd.connect_to_pc() return _json(await self.__msd.connect_to_pc())
state = self.__msd.get_state()
await self.__broadcast_event(_Events.MSD_STATE, state)
else: else:
raise BadRequest("Invalid param 'to'") raise BadRequest("Invalid param 'to'")
return _json(state)
@_wrap_exceptions_for_web("Can't write data to mass-storage device") @_wrap_exceptions_for_web("Can't write data to mass-storage device")
async def __msd_write_handler(self, request: aiohttp.web.Request) -> aiohttp.web.Response: async def __msd_write_handler(self, request: aiohttp.web.Request) -> aiohttp.web.Response:
@ -334,17 +328,16 @@ class Server: # pylint: disable=too-many-instance-attributes
reader = await request.multipart() reader = await request.multipart()
written = 0 written = 0
try: try:
field = await reader.next()
if not field or field.name != "image_name":
raise BadRequest("Missing 'image_name' field")
image_name = (await field.read()).decode("utf-8")[:256]
field = await reader.next()
if not field or field.name != "image_data":
raise BadRequest("Missing 'image_data' field")
async with self.__msd: async with self.__msd:
await self.__broadcast_event(_Events.MSD_STATE, self.__msd.get_state()) field = await reader.next()
if not field or field.name != "image_name":
raise BadRequest("Missing 'image_name' field")
image_name = (await field.read()).decode("utf-8")[:256]
field = await reader.next()
if not field or field.name != "image_data":
raise BadRequest("Missing 'image_data' field")
logger.info("Writing image %r to mass-storage device ...", image_name) logger.info("Writing image %r to mass-storage device ...", image_name)
await self.__msd.write_image_info(image_name, False) await self.__msd.write_image_info(image_name, False)
while True: while True:
@ -354,7 +347,6 @@ class Server: # pylint: disable=too-many-instance-attributes
written = await self.__msd.write_image_chunk(chunk) written = await self.__msd.write_image_chunk(chunk)
await self.__msd.write_image_info(image_name, True) await self.__msd.write_image_info(image_name, True)
finally: finally:
await self.__broadcast_event(_Events.MSD_STATE, self.__msd.get_state())
if written != 0: if written != 0:
logger.info("Written %d bytes to mass-storage device", written) logger.info("Written %d bytes to mass-storage device", written)
return _json({"written": written}) return _json({"written": written})
@ -450,27 +442,31 @@ class Server: # pylint: disable=too-many-instance-attributes
@_system_task @_system_task
async def __poll_atx_state(self) -> None: async def __poll_atx_state(self) -> None:
async for state in self.__atx.poll_state(): async for state in self.__atx.poll_state():
if self.__sockets: await self.__broadcast_event(_Events.ATX_STATE, state)
await self.__broadcast_event(_Events.ATX_STATE, state)
@_system_task
async def __poll_msd_state(self) -> None:
async for state in self.__msd.poll_state():
await self.__broadcast_event(_Events.MSD_STATE, state)
@_system_task @_system_task
async def __poll_streamer_state(self) -> None: async def __poll_streamer_state(self) -> None:
async for state in self.__streamer.poll_state(): async for state in self.__streamer.poll_state():
if self.__sockets: await self.__broadcast_event(_Events.STREAMER_STATE, state)
await self.__broadcast_event(_Events.STREAMER_STATE, state)
async def __broadcast_event(self, event_type: _Events, event_attrs: Dict) -> None: async def __broadcast_event(self, event_type: _Events, event_attrs: Dict) -> None:
await asyncio.gather(*[ if self.__sockets:
ws.send_str(json.dumps({ await asyncio.gather(*[
"msg_type": "event", ws.send_str(json.dumps({
"msg": { "msg_type": "event",
"event": event_type.value, "msg": {
"event_attrs": event_attrs, "event": event_type.value,
}, "event_attrs": event_attrs,
})) },
for ws in list(self.__sockets) }))
if not ws.closed and ws._req.transport # pylint: disable=protected-access for ws in list(self.__sockets)
], return_exceptions=True) if not ws.closed and ws._req.transport # pylint: disable=protected-access
], return_exceptions=True)
async def __register_socket(self, ws: aiohttp.web.WebSocketResponse) -> None: async def __register_socket(self, ws: aiohttp.web.WebSocketResponse) -> None:
async with self.__sockets_lock: async with self.__sockets_lock: