summaryrefslogtreecommitdiff
path: root/kvmd/kvmd/server.py
blob: 81f928e487b08158721947e5ac5ffdb93df6f433 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
import os
import signal
import asyncio
import time

from typing import List
from typing import Set
from typing import Callable
from typing import Optional

import aiohttp

from .atx import Atx
from .streamer import Streamer
from .ps2 import Ps2Keyboard

from .logging import get_logger


# =====
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


class Server:  # pylint: disable=too-many-instance-attributes
    def __init__(
        self,
        atx: Atx,
        streamer: Streamer,
        keyboard: Ps2Keyboard,
        heartbeat: float,
        atx_leds_poll: float,
        video_shutdown_delay: float,
        loop: asyncio.AbstractEventLoop,
    ) -> None:

        self.__atx = atx
        self.__streamer = streamer
        self.__heartbeat = heartbeat
        self.__keyboard = keyboard
        self.__video_shutdown_delay = video_shutdown_delay
        self.__atx_leds_poll = atx_leds_poll
        self.__loop = loop

        self.__sockets: Set[aiohttp.web.WebSocketResponse] = set()
        self.__sockets_lock = asyncio.Lock()

        self.__system_tasks: List[asyncio.Task] = []

    def run(self, host: str, port: int) -> None:
        self.__keyboard.start()

        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_tasks.extend([
            self.__loop.create_task(self.__monitor_keyboard()),
            self.__loop.create_task(self.__stream_controller()),
            self.__loop.create_task(self.__poll_dead_sockets()),
            self.__loop.create_task(self.__poll_atx_leds()),
        ])

        aiohttp.web.run_app(
            app=app,
            host=host,
            port=port,
            print=(lambda text: [get_logger().info(line.strip()) for line in text.strip().splitlines()]),  # type: ignore
        )

    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(heartbeat=self.__heartbeat)
        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:
        logger = get_logger(0)

        logger.info("Cancelling system tasks ...")
        for task in self.__system_tasks:
            task.cancel()
        await asyncio.gather(*self.__system_tasks)

        logger.info("Disconnecting clients ...")
        for ws in list(self.__sockets):
            await self.__remove_socket(ws)

    async def __on_cleanup(self, _: aiohttp.web.Application) -> None:
        if self.__keyboard.is_alive():
            self.__keyboard.stop()
        if self.__streamer.is_running():
            await self.__streamer.stop()

    @_system_task
    async def __monitor_keyboard(self) -> None:
        while self.__keyboard.is_alive():
            await asyncio.sleep(0.1)
        raise RuntimeError("Keyboard dead")

    @_system_task
    async def __stream_controller(self) -> None:
        prev = 0
        shutdown_at = 0.0
        while True:
            cur = len(self.__sockets)
            if prev == 0 and cur > 0:
                if not self.__streamer.is_running():
                    await self.__streamer.start()
            elif prev > 0 and cur == 0:
                shutdown_at = time.time() + self.__video_shutdown_delay
            elif prev == 0 and cur == 0 and time.time() > shutdown_at:
                if self.__streamer.is_running():
                    await self.__streamer.stop()
            prev = cur
            await asyncio.sleep(0.1)

    @_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.__atx_leds_poll)

    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 an 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

    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