mirror of
https://github.com/mofeng-git/One-KVM.git
synced 2026-02-02 02:51:53 +08:00
initial commit
This commit is contained in:
Devaev Maxim
197
kvmd/kvmd/__init__.py
Normal file
197
kvmd/kvmd/__init__.py
Normal file
@@ -0,0 +1,197 @@
|
||||
import asyncio
|
||||
import argparse
|
||||
import logging
|
||||
import logging.config
|
||||
|
||||
from typing import List
|
||||
from typing import Dict
|
||||
from typing import Set
|
||||
from typing import Callable
|
||||
from typing import Optional
|
||||
|
||||
from contextlog import get_logger
|
||||
from contextlog import patch_logging
|
||||
from contextlog import patch_threading
|
||||
|
||||
from RPi import GPIO
|
||||
|
||||
import aiohttp
|
||||
|
||||
import yaml
|
||||
|
||||
from .atx import Atx
|
||||
from .streamer import Streamer
|
||||
|
||||
|
||||
# =====
|
||||
def _system_task(method: Callable) -> Callable:
|
||||
async def wrap(self: "_Application") -> None:
|
||||
try:
|
||||
await method(self)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
get_logger().exception("Unhandled exception")
|
||||
raise SystemExit(1)
|
||||
return wrap
|
||||
|
||||
|
||||
class _Application:
|
||||
def __init__(self, config: Dict) -> None:
|
||||
self.__config = config
|
||||
|
||||
self.__loop = asyncio.get_event_loop()
|
||||
self.__sockets: Set[aiohttp.web.WebSocketResponse] = set()
|
||||
self.__sockets_lock = asyncio.Lock()
|
||||
|
||||
GPIO.setmode(GPIO.BCM)
|
||||
|
||||
self.__atx = Atx(
|
||||
power_led=self.__config["atx"]["leds"]["pinout"]["power"],
|
||||
hdd_led=self.__config["atx"]["leds"]["pinout"]["hdd"],
|
||||
power_switch=self.__config["atx"]["switches"]["pinout"]["power"],
|
||||
reset_switch=self.__config["atx"]["switches"]["pinout"]["reset"],
|
||||
click_delay=self.__config["atx"]["switches"]["click_delay"],
|
||||
long_click_delay=self.__config["atx"]["switches"]["long_click_delay"],
|
||||
)
|
||||
|
||||
self.__streamer = Streamer(
|
||||
cap_power=self.__config["video"]["pinout"]["cap"],
|
||||
vga_power=self.__config["video"]["pinout"]["vga"],
|
||||
sync_delay=self.__config["video"]["sync_delay"],
|
||||
mjpg_streamer=self.__config["video"]["mjpg_streamer"],
|
||||
)
|
||||
|
||||
self.__system_futures: List[asyncio.Future] = []
|
||||
|
||||
def run(self) -> None:
|
||||
app = aiohttp.web.Application(loop=self.__loop)
|
||||
app.router.add_get("/", self.__root_handler)
|
||||
app.router.add_get("/ws", self.__ws_handler)
|
||||
app.on_shutdown.append(self.__on_shutdown)
|
||||
app.on_cleanup.append(self.__on_cleanup)
|
||||
|
||||
self.__system_futures.extend([
|
||||
asyncio.ensure_future(self.__poll_dead_sockets(), loop=self.__loop),
|
||||
asyncio.ensure_future(self.__poll_atx_leds(), loop=self.__loop),
|
||||
asyncio.ensure_future(self.__poll_streamer_events(), loop=self.__loop),
|
||||
])
|
||||
|
||||
aiohttp.web.run_app(
|
||||
app=app,
|
||||
host=self.__config["server"]["host"],
|
||||
port=self.__config["server"]["port"],
|
||||
print=(lambda text: [get_logger().info(line.strip()) for line in text.strip().splitlines()]),
|
||||
)
|
||||
|
||||
async def __root_handler(self, _: aiohttp.web.Request) -> aiohttp.web.Response:
|
||||
return aiohttp.web.Response(text="OK")
|
||||
|
||||
async def __ws_handler(self, request: aiohttp.web.Request) -> aiohttp.web.WebSocketResponse:
|
||||
ws = aiohttp.web.WebSocketResponse(**self.__config["ws"])
|
||||
await ws.prepare(request)
|
||||
await self.__register_socket(ws)
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.web.WSMsgType.TEXT:
|
||||
retval = await self.__execute_command(msg.data)
|
||||
if retval:
|
||||
await ws.send_str(retval)
|
||||
else:
|
||||
break
|
||||
return ws
|
||||
|
||||
async def __on_shutdown(self, _: aiohttp.web.Application) -> None:
|
||||
get_logger().info("Shutting down ...")
|
||||
for ws in list(self.__sockets):
|
||||
await self.__remove_socket(ws)
|
||||
|
||||
async def __on_cleanup(self, _: aiohttp.web.Application) -> None:
|
||||
logger = get_logger()
|
||||
|
||||
logger.info("Cancelling tasks ...")
|
||||
for future in self.__system_futures:
|
||||
future.cancel()
|
||||
await asyncio.gather(*self.__system_futures)
|
||||
|
||||
logger.info("Cleaning up GPIO ...")
|
||||
GPIO.cleanup()
|
||||
|
||||
logger.info("Bye-bye")
|
||||
|
||||
@_system_task
|
||||
async def __poll_dead_sockets(self) -> None:
|
||||
while True:
|
||||
for ws in list(self.__sockets):
|
||||
if ws.closed or not ws._req.transport: # pylint: disable=protected-access
|
||||
await self.__remove_socket(ws)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
@_system_task
|
||||
async def __poll_atx_leds(self) -> None:
|
||||
while True:
|
||||
if self.__sockets:
|
||||
await self.__broadcast("EVENT atx_leds %d %d" % (self.__atx.get_leds()))
|
||||
await asyncio.sleep(self.__config["atx"]["leds"]["poll"])
|
||||
|
||||
@_system_task
|
||||
async def __poll_streamer_events(self) -> None:
|
||||
async for event in self.__streamer.events():
|
||||
await self.__broadcast("EVENT %s" % (event))
|
||||
|
||||
async def __broadcast(self, msg: str) -> None:
|
||||
await asyncio.gather(*[
|
||||
ws.send_str(msg)
|
||||
for ws in list(self.__sockets)
|
||||
if not ws.closed and ws._req.transport # pylint: disable=protected-access
|
||||
], return_exceptions=True)
|
||||
|
||||
async def __execute_command(self, command: str) -> Optional[str]:
|
||||
(command, args) = (command.strip().split(" ", maxsplit=1) + [""])[:2]
|
||||
if command == "CLICK":
|
||||
method = {
|
||||
"power": self.__atx.click_power,
|
||||
"power_long": self.__atx.click_power_long,
|
||||
"reset": self.__atx.click_reset,
|
||||
}.get(args)
|
||||
if method:
|
||||
await method()
|
||||
return None
|
||||
get_logger().warning("Received incorrect command: %r", command)
|
||||
return "ERROR incorrect command"
|
||||
|
||||
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
|
||||
if len(self.__sockets) == 1:
|
||||
await self.__streamer.start()
|
||||
|
||||
async def __remove_socket(self, ws: aiohttp.web.WebSocketResponse) -> None:
|
||||
async with self.__sockets_lock:
|
||||
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
|
||||
if not self.__sockets:
|
||||
await self.__streamer.stop()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
patch_logging()
|
||||
patch_threading()
|
||||
get_logger(app="kvmd")
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-c", "--config", default="kvmd.yaml", metavar="<path>")
|
||||
options = parser.parse_args()
|
||||
|
||||
with open(options.config) as config_file:
|
||||
config = yaml.load(config_file)
|
||||
logging.captureWarnings(True)
|
||||
logging.config.dictConfig(config["logging"])
|
||||
|
||||
_Application(config["kvmd"]).run()
|
||||
2
kvmd/kvmd/__main__.py
Normal file
2
kvmd/kvmd/__main__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
from . import main
|
||||
main()
|
||||
62
kvmd/kvmd/atx.py
Normal file
62
kvmd/kvmd/atx.py
Normal file
@@ -0,0 +1,62 @@
|
||||
import asyncio
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
from contextlog import get_logger
|
||||
|
||||
from RPi import GPIO
|
||||
|
||||
|
||||
# =====
|
||||
class Atx:
|
||||
def __init__(
|
||||
self,
|
||||
power_led: int,
|
||||
hdd_led: int,
|
||||
power_switch: int,
|
||||
reset_switch: int,
|
||||
click_delay: float,
|
||||
long_click_delay: float,
|
||||
) -> None:
|
||||
|
||||
self.__power_led = self.__set_output_pin(power_led)
|
||||
self.__hdd_led = self.__set_output_pin(hdd_led)
|
||||
|
||||
self.__power_switch = self.__set_output_pin(power_switch)
|
||||
self.__reset_switch = self.__set_output_pin(reset_switch)
|
||||
self.__click_delay = click_delay
|
||||
self.__long_click_delay = long_click_delay
|
||||
|
||||
self.__lock = asyncio.Lock()
|
||||
|
||||
def __set_output_pin(self, pin: int) -> int:
|
||||
GPIO.setup(pin, GPIO.OUT)
|
||||
GPIO.output(pin, False)
|
||||
return pin
|
||||
|
||||
def get_leds(self) -> Tuple[bool, bool]:
|
||||
return (
|
||||
not GPIO.input(self.__power_led),
|
||||
not GPIO.input(self.__hdd_led),
|
||||
)
|
||||
|
||||
async def click_power(self) -> None:
|
||||
if (await self.__click(self.__power_switch, self.__click_delay)):
|
||||
get_logger().info("Clicked power")
|
||||
|
||||
async def click_power_long(self) -> None:
|
||||
if (await self.__click(self.__power_switch, self.__long_click_delay)):
|
||||
get_logger().info("Clicked power (long press)")
|
||||
|
||||
async def click_reset(self) -> None:
|
||||
if (await self.__click(self.__reset_switch, self.__click_delay)):
|
||||
get_logger().info("Clicked reset")
|
||||
|
||||
async def __click(self, pin: int, delay: float) -> bool:
|
||||
if not self.__lock.locked():
|
||||
async with self.__lock:
|
||||
for flag in (True, False):
|
||||
GPIO.output(pin, flag)
|
||||
await asyncio.sleep(delay)
|
||||
return True
|
||||
return False
|
||||
120
kvmd/kvmd/streamer.py
Normal file
120
kvmd/kvmd/streamer.py
Normal file
@@ -0,0 +1,120 @@
|
||||
import asyncio
|
||||
import asyncio.subprocess
|
||||
|
||||
from typing import Dict
|
||||
from typing import AsyncIterator
|
||||
from typing import Optional
|
||||
|
||||
from contextlog import get_logger
|
||||
|
||||
from RPi import GPIO
|
||||
|
||||
|
||||
# =====
|
||||
class Streamer:
|
||||
def __init__(
|
||||
self,
|
||||
cap_power: int,
|
||||
vga_power: int,
|
||||
sync_delay: float,
|
||||
mjpg_streamer: Dict,
|
||||
) -> None:
|
||||
|
||||
self.__cap_power = self.__set_output_pin(cap_power)
|
||||
self.__vga_power = self.__set_output_pin(vga_power)
|
||||
self.__sync_delay = sync_delay
|
||||
|
||||
self.__cmd = (
|
||||
"%(prog)s"
|
||||
" -i 'input_uvc.so -d %(device)s -e %(every)s -y -n -r %(width)sx%(height)s'"
|
||||
" -o 'output_http.so -p -l %(host)s %(port)s'"
|
||||
) % (mjpg_streamer)
|
||||
|
||||
self.__lock = asyncio.Lock()
|
||||
self.__events_queue: asyncio.Queue = asyncio.Queue()
|
||||
self.__proc_future: Optional[asyncio.Future] = None
|
||||
|
||||
def __set_output_pin(self, pin: int) -> int:
|
||||
GPIO.setup(pin, GPIO.OUT)
|
||||
GPIO.output(pin, False)
|
||||
return pin
|
||||
|
||||
async def events(self) -> AsyncIterator[str]:
|
||||
while True:
|
||||
yield (await self.__events_queue.get())
|
||||
|
||||
async def start(self) -> None:
|
||||
async with self.__lock:
|
||||
get_logger().info("Starting mjpg_streamer ...")
|
||||
assert not self.__proc_future
|
||||
await self.__set_hw_enabled(True)
|
||||
self.__proc_future = asyncio.ensure_future(self.__process(), loop=asyncio.get_event_loop())
|
||||
|
||||
async def stop(self) -> None:
|
||||
async with self.__lock:
|
||||
get_logger().info("Stopping mjpg_streamer ...")
|
||||
if self.__proc_future:
|
||||
self.__proc_future.cancel()
|
||||
await asyncio.gather(self.__proc_future, return_exceptions=True)
|
||||
await self.__set_hw_enabled(False)
|
||||
self.__proc_future = None
|
||||
await self.__events_queue.put("mjpg_streamer stopped")
|
||||
|
||||
async def __set_hw_enabled(self, enabled: bool) -> None:
|
||||
# This sequence is important for enable
|
||||
GPIO.output(self.__cap_power, enabled)
|
||||
if enabled:
|
||||
await asyncio.sleep(self.__sync_delay)
|
||||
GPIO.output(self.__vga_power, enabled)
|
||||
await asyncio.sleep(self.__sync_delay)
|
||||
|
||||
async def __process(self) -> None:
|
||||
logger = get_logger()
|
||||
|
||||
proc: Optional[asyncio.subprocess.Process] = None # pylint: disable=no-member
|
||||
while True:
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_shell(
|
||||
self.__cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.STDOUT,
|
||||
)
|
||||
|
||||
logger.info("Started mjpg_streamer pid=%d: %s", proc.pid, self.__cmd)
|
||||
await self.__events_queue.put("mjpg_streamer started")
|
||||
|
||||
empty = 0
|
||||
while proc.returncode is None:
|
||||
line = (await proc.stdout.readline()).decode(errors="ignore").strip()
|
||||
if line:
|
||||
logger.info("mjpg_streamer: %s", line)
|
||||
empty = 0
|
||||
else:
|
||||
empty += 1
|
||||
if empty == 100: # asyncio bug
|
||||
break
|
||||
|
||||
await self.__kill(proc)
|
||||
raise RuntimeError("WTF")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as err:
|
||||
if proc:
|
||||
logger.error("Unexpected finished mjpg_streamer pid=%d with retcode=%d", proc.pid, proc.returncode)
|
||||
else:
|
||||
logger.error("Can't start mjpg_streamer: %s", err)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
if proc:
|
||||
await self.__kill(proc)
|
||||
|
||||
async def __kill(self, proc: asyncio.subprocess.Process) -> None: # pylint: disable=no-member
|
||||
try:
|
||||
proc.terminate()
|
||||
await asyncio.sleep(1)
|
||||
if proc.returncode is None:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
except Exception:
|
||||
pass
|
||||
Reference in New Issue
Block a user