70 lines
1.9 KiB
Python
70 lines
1.9 KiB
Python
import asyncio
|
|
from uuid import UUID
|
|
|
|
from tinydb import TinyDB
|
|
|
|
from settings import settings
|
|
from proto import Computer, DynamicComputerState
|
|
|
|
db = TinyDB(settings.database_path)
|
|
computers = db.table("computers")
|
|
|
|
|
|
class StateManager:
|
|
def __init__(self):
|
|
self.websockets = set()
|
|
self.connected_computers = set()
|
|
|
|
self.current_state = None
|
|
self.update_state()
|
|
|
|
def update_state(self):
|
|
|
|
self.current_state = {"computers": []}
|
|
|
|
for computer in computers.all():
|
|
|
|
static = Computer.parse_obj(computer)
|
|
|
|
dynamic = DynamicComputerState(
|
|
is_online=static.uuid in self.connected_computers,
|
|
)
|
|
|
|
self.current_state["computers"].append(
|
|
{"static": static.dict(), "dynamic": dynamic.dict()}
|
|
)
|
|
|
|
async def push_state(self, socket):
|
|
try:
|
|
await socket.send_json(self.current_state)
|
|
except RuntimeError:
|
|
print("dead socket?")
|
|
self.websockets.remove(socket)
|
|
|
|
async def on_connect(self, socket):
|
|
self.websockets.add(socket)
|
|
await self.push_state(socket)
|
|
|
|
async def on_disconnect(self, socket):
|
|
self.websockets.remove(socket)
|
|
|
|
async def on_change(self):
|
|
self.update_state()
|
|
await asyncio.gather(*[self.push_state(socket) for socket in self.websockets])
|
|
|
|
async def on_computer_register(self, computer: Computer):
|
|
|
|
computers.insert(computer.dict())
|
|
await self.on_change()
|
|
|
|
async def on_computer_connect(self, computer_id: UUID):
|
|
self.connected_computers.add(computer_id)
|
|
await self.on_change()
|
|
|
|
async def on_computer_disconnect(self, computer_id: UUID):
|
|
self.connected_computers.remove(computer_id)
|
|
await self.on_change()
|
|
|
|
async def on_computer_message(self, computer_id: UUID, message):
|
|
print(f"[on_computer_message] UUID: {computer_id} Message: {message!r}")
|