GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2026-01-29 00:51:53 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactoring

Browse Source
This commit is contained in:
Devaev Maxim
2018-12-13 20:39:01 +03:00
parent 972c4928cc
commit 1c64b7b0e8
2 changed files with 108 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2,6 +2,7 @@ import os
import signal import signal
import socket import socket
import asyncio import asyncio
import inspect
import json import json
import time import time
@@ -32,18 +33,6 @@ from .streamer import Streamer
# ===== # =====
def _system_task(method: Callable) -> Callable:
async def wrap(self: "Server") -> None:
try:
await method(self)
except asyncio.CancelledError:
pass
except Exception:
get_logger().exception("Unhandled exception, killing myself ...")
os.kill(os.getpid(), signal.SIGTERM)
return wrap
def _json(result: Optional[Dict]=None, status: int=200) -> aiohttp.web.Response: def _json(result: Optional[Dict]=None, status: int=200) -> aiohttp.web.Response:
return aiohttp.web.Response( return aiohttp.web.Response(
text=json.dumps({ text=json.dumps({
@@ -55,26 +44,64 @@ def _json(result: Optional[Dict]=None, status: int=200) -> aiohttp.web.Response:
) )
def _json_exception(msg: str, err: Exception, status: int) -> aiohttp.web.Response: def _json_exception(err: Exception, status: int) -> aiohttp.web.Response:
msg = "%s: %s" % (msg, err) name = type(err).__name__
get_logger().error(msg) msg = str(err)
get_logger().error("API error: %s: %s", name, msg)
return _json({ return _json({
"error": type(err).__name__, "error": name,
"error_msg": msg, "error_msg": msg,
}, status=status) }, status=status)
class BadRequest(Exception): class BadRequestError(Exception):
pass pass
_ATTR_EXPOSED = "exposed"
_ATTR_EXPOSED_METHOD = "exposed_method"
_ATTR_EXPOSED_PATH = "exposed_path"
_ATTR_SYSTEM_TASK = "system_task"
def _exposed(http_method: str, path: str) -> Callable:
def make_wrapper(method: Callable) -> Callable:
async def wrap(self: "Server", request: aiohttp.web.Request) -> aiohttp.web.Response:
try:
return (await method(self, request))
except RegionIsBusyError as err:
return _json_exception(err, 409)
except (BadRequestError, MsdOperationError) as err:
return _json_exception(err, 400)
setattr(wrap, _ATTR_EXPOSED, True)
setattr(wrap, _ATTR_EXPOSED_METHOD, http_method)
setattr(wrap, _ATTR_EXPOSED_PATH, path)
return wrap
return make_wrapper
def _system_task(method: Callable) -> Callable:
async def wrap(self: "Server") -> None:
try:
await method(self)
except asyncio.CancelledError:
pass
except Exception:
get_logger().exception("Unhandled exception, killing myself ...")
os.kill(os.getpid(), signal.SIGTERM)
setattr(wrap, _ATTR_SYSTEM_TASK, True)
return wrap
def _valid_bool(name: str, flag: Optional[str]) -> bool: def _valid_bool(name: str, flag: Optional[str]) -> bool:
flag = str(flag).strip().lower() flag = str(flag).strip().lower()
if flag in ["1", "true", "yes"]: if flag in ["1", "true", "yes"]:
return True return True
elif flag in ["0", "false", "no"]: elif flag in ["0", "false", "no"]:
return False return False
raise BadRequest("Invalid param '%s'" % (name)) raise BadRequestError("Invalid param '%s'" % (name))
def _valid_int(name: str, value: Optional[str], min_value: Optional[int]=None, max_value: Optional[int]=None) -> int: def _valid_int(name: str, value: Optional[str], min_value: Optional[int]=None, max_value: Optional[int]=None) -> int:
@@ -87,20 +114,7 @@ def _valid_int(name: str, value: Optional[str], min_value: Optional[int]=None, m
raise ValueError() raise ValueError()
return value_int return value_int
except Exception: except Exception:
raise BadRequest("Invalid param %r" % (name)) raise BadRequestError("Invalid param %r" % (name))
def _wrap_exceptions_for_web(msg: str) -> Callable:
def make_wrapper(method: Callable) -> Callable:
async def wrap(self: "Server", request: aiohttp.web.Request) -> aiohttp.web.Response:
try:
return (await method(self, request))
except RegionIsBusyError as err:
return _json_exception(msg, err, 409)
except (BadRequest, MsdOperationError) as err:
return _json_exception(msg, err, 400)
return wrap
return make_wrapper
class _Events(Enum): class _Events(Enum):
@@ -156,37 +170,20 @@ class Server: # pylint: disable=too-many-instance-attributes
setproctitle.setproctitle("[main] " + setproctitle.getproctitle()) setproctitle.setproctitle("[main] " + setproctitle.getproctitle())
app = aiohttp.web.Application(loop=self.__loop) app = aiohttp.web.Application(loop=self.__loop)
app.router.add_get("/info", self.__info_handler)
app.router.add_get("/log", self.__log_handler)
app.router.add_get("/ws", self.__ws_handler)
app.router.add_post("/hid/reset", self.__hid_reset_handler)
app.router.add_get("/atx", self.__atx_state_handler)
app.router.add_post("/atx/click", self.__atx_click_handler)
app.router.add_get("/msd", self.__msd_state_handler)
app.router.add_post("/msd/connect", self.__msd_connect_handler)
app.router.add_post("/msd/write", self.__msd_write_handler)
app.router.add_post("/msd/reset", self.__msd_reset_handler)
app.router.add_get("/streamer", self.__streamer_state_handler)
app.router.add_post("/streamer/set_params", self.__streamer_set_params_handler)
app.router.add_post("/streamer/reset", self.__streamer_reset_handler)
app.on_shutdown.append(self.__on_shutdown) app.on_shutdown.append(self.__on_shutdown)
app.on_cleanup.append(self.__on_cleanup) app.on_cleanup.append(self.__on_cleanup)
self.__system_tasks.extend([ for name in dir(self):
self.__loop.create_task(self.__hid_watchdog()), method = getattr(self, name)
self.__loop.create_task(self.__stream_controller()), if inspect.ismethod(method):
self.__loop.create_task(self.__poll_dead_sockets()), if getattr(method, _ATTR_SYSTEM_TASK, False):
self.__loop.create_task(self.__poll_atx_state()), self.__system_tasks.append(self.__loop.create_task(method()))
self.__loop.create_task(self.__poll_msd_state()), elif getattr(method, _ATTR_EXPOSED, False):
self.__loop.create_task(self.__poll_streamer_state()), app.router.add_route(
]) getattr(method, _ATTR_EXPOSED_METHOD),
getattr(method, _ATTR_EXPOSED_PATH),
method,
)
assert port or unix_path assert port or unix_path
if unix_path: if unix_path:
@@ -215,10 +212,11 @@ class Server: # pylint: disable=too-many-instance-attributes
# ===== SYSTEM # ===== SYSTEM
@_exposed("GET", "/info")
async def __info_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response: async def __info_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response:
return _json(await self.__make_info()) return _json(await self.__make_info())
@_wrap_exceptions_for_web("Log error") @_exposed("GET", "/log")
async def __log_handler(self, request: aiohttp.web.Request) -> aiohttp.web.StreamResponse: async def __log_handler(self, request: aiohttp.web.Request) -> aiohttp.web.StreamResponse:
seek = _valid_int("seek", request.query.get("seek", "0"), 0) seek = _valid_int("seek", request.query.get("seek", "0"), 0)
follow = _valid_bool("follow", request.query.get("follow", "false")) follow = _valid_bool("follow", request.query.get("follow", "false"))
@@ -234,6 +232,7 @@ class Server: # pylint: disable=too-many-instance-attributes
# ===== WEBSOCKET # ===== WEBSOCKET
@_exposed("GET", "/ws")
async def __ws_handler(self, request: aiohttp.web.Request) -> aiohttp.web.WebSocketResponse: async def __ws_handler(self, request: aiohttp.web.Request) -> aiohttp.web.WebSocketResponse:
logger = get_logger(0) logger = get_logger(0)
ws = aiohttp.web.WebSocketResponse(heartbeat=self.__heartbeat) ws = aiohttp.web.WebSocketResponse(heartbeat=self.__heartbeat)
@@ -298,16 +297,18 @@ class Server: # pylint: disable=too-many-instance-attributes
# ===== HID # ===== HID
@_exposed("POST", "/hid/reset")
async def __hid_reset_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response: async def __hid_reset_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response:
await self.__hid.reset() await self.__hid.reset()
return _json() return _json()
# ===== ATX # ===== ATX
@_exposed("GET", "/atx")
async def __atx_state_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response: async def __atx_state_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response:
return _json(self.__atx.get_state()) return _json(self.__atx.get_state())
@_wrap_exceptions_for_web("Click error") @_exposed("POST", "/atx/click")
async def __atx_click_handler(self, request: aiohttp.web.Request) -> aiohttp.web.Response: async def __atx_click_handler(self, request: aiohttp.web.Request) -> aiohttp.web.Response:
button = request.query.get("button") button = request.query.get("button")
clicker = { clicker = {
@@ -316,16 +317,17 @@ class Server: # pylint: disable=too-many-instance-attributes
"reset": self.__atx.click_reset, "reset": self.__atx.click_reset,
}.get(button) }.get(button)
if not clicker: if not clicker:
raise BadRequest("Invalid param 'button'") raise BadRequestError("Invalid param 'button'")
await clicker() await clicker()
return _json({"clicked": button}) return _json({"clicked": button})
# ===== MSD # ===== MSD
@_exposed("GET", "/msd")
async def __msd_state_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response: async def __msd_state_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response:
return _json(self.__msd.get_state()) return _json(self.__msd.get_state())
@_wrap_exceptions_for_web("Mass-storage error") @_exposed("POST", "/msd/connect")
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":
@@ -333,9 +335,9 @@ class Server: # pylint: disable=too-many-instance-attributes
elif to == "server": elif to == "server":
return _json(await self.__msd.connect_to_pc()) return _json(await self.__msd.connect_to_pc())
else: else:
raise BadRequest("Invalid param 'to'") raise BadRequestError("Invalid param 'to'")
@_wrap_exceptions_for_web("Can't write data to mass-storage device") @_exposed("POST", "/msd/write")
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:
logger = get_logger(0) logger = get_logger(0)
reader = await request.multipart() reader = await request.multipart()
@@ -344,12 +346,12 @@ class Server: # pylint: disable=too-many-instance-attributes
async with self.__msd: async with self.__msd:
field = await reader.next() field = await reader.next()
if not field or field.name != "image_name": if not field or field.name != "image_name":
raise BadRequest("Missing 'image_name' field") raise BadRequestError("Missing 'image_name' field")
image_name = (await field.read()).decode("utf-8")[:256] image_name = (await field.read()).decode("utf-8")[:256]
field = await reader.next() field = await reader.next()
if not field or field.name != "image_data": if not field or field.name != "image_data":
raise BadRequest("Missing 'image_data' field") raise BadRequestError("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)
@@ -364,17 +366,18 @@ class Server: # pylint: disable=too-many-instance-attributes
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})
@_wrap_exceptions_for_web("Mass-storage error") @_exposed("POST", "/msd/reset")
async def __msd_reset_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response: async def __msd_reset_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response:
await self.__msd.reset() await self.__msd.reset()
return _json() return _json()
# ===== STREAMER # ===== STREAMER
@_exposed("GET", "/streamer")
async def __streamer_state_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response: async def __streamer_state_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response:
return _json(await self.__streamer.get_state()) return _json(await self.__streamer.get_state())
@_wrap_exceptions_for_web("Can't set stream params") @_exposed("POST", "/streamer/set_params")
async def __streamer_set_params_handler(self, request: aiohttp.web.Request) -> aiohttp.web.Response: async def __streamer_set_params_handler(self, request: aiohttp.web.Request) -> aiohttp.web.Response:
for (name, validator) in [ for (name, validator) in [
("quality", lambda arg: _valid_int("quality", arg, 1, 100)), ("quality", lambda arg: _valid_int("quality", arg, 1, 100)),
@@ -385,6 +388,7 @@ class Server: # pylint: disable=too-many-instance-attributes
self.__streamer_params[name] = validator(value) self.__streamer_params[name] = validator(value)
return _json() return _json()
@_exposed("POST", "/streamer/reset")
async def __streamer_reset_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response: async def __streamer_reset_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response:
self.__reset_streamer = True self.__reset_streamer = True
return _json() return _json()
@@ -413,6 +417,39 @@ class Server: # pylint: disable=too-many-instance-attributes
await self.__msd.cleanup() await self.__msd.cleanup()
await self.__hid.cleanup() await self.__hid.cleanup()
async def __broadcast_event(self, event_type: _Events, event_attrs: Dict) -> None:
if self.__sockets:
await asyncio.gather(*[
ws.send_str(json.dumps({
"msg_type": "event",
"msg": {
"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
], return_exceptions=True)
async def __register_socket(self, ws: aiohttp.web.WebSocketResponse) -> None:
async with self.__sockets_lock:
self.__sockets.add(ws)
get_logger().info("Registered new client socket: remote=%s; id=%d; active=%d",
ws._req.remote, id(ws), len(self.__sockets)) # pylint: disable=protected-access
async def __remove_socket(self, ws: aiohttp.web.WebSocketResponse) -> None:
async with self.__sockets_lock:
await self.__hid.clear_events()
try:
self.__sockets.remove(ws)
get_logger().info("Removed client socket: remote=%s; id=%d; active=%d",
ws._req.remote, id(ws), len(self.__sockets)) # pylint: disable=protected-access
await ws.close()
except Exception:
pass
# ===== SYSTEM TASKS
@_system_task @_system_task
async def __hid_watchdog(self) -> None: async def __hid_watchdog(self) -> None:
while self.__hid.is_alive(): while self.__hid.is_alive():
@@ -466,34 +503,3 @@ class Server: # pylint: disable=too-many-instance-attributes
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():
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:
if self.__sockets:
await asyncio.gather(*[
ws.send_str(json.dumps({
"msg_type": "event",
"msg": {
"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
], return_exceptions=True)
async def __register_socket(self, ws: aiohttp.web.WebSocketResponse) -> None:
async with self.__sockets_lock:
self.__sockets.add(ws)
get_logger().info("Registered new client socket: remote=%s; id=%d; active=%d",
ws._req.remote, id(ws), len(self.__sockets)) # pylint: disable=protected-access
async def __remove_socket(self, ws: aiohttp.web.WebSocketResponse) -> None:
async with self.__sockets_lock:
await self.__hid.clear_events()
try:
self.__sockets.remove(ws)
get_logger().info("Removed client socket: remote=%s; id=%d; active=%d",
ws._req.remote, id(ws), len(self.__sockets)) # pylint: disable=protected-access
await ws.close()
except Exception:
pass

2
testenv/tox.ini
View File

@@ -26,7 +26,7 @@ deps =
-rrequirements.txt -rrequirements.txt
[testenv:vulture] [testenv:vulture]
commands = vulture kvmd genmap.py testenv/vulture-wl.py commands = vulture --ignore-decorators=@_exposed,@_system_task kvmd genmap.py testenv/vulture-wl.py
deps = deps =
vulture vulture
-rrequirements.txt -rrequirements.txt