Update proto to include dynamic information

backend/auth.py

@@ -1,6 +1,10 @@
+from typing import *
+
+from pydantic import ValidationError
 from jose import jwt
 
 from settings import settings
+from proto import FrontendToken, ComputerToken
 
 
 def encode(data):
@@ -11,23 +15,21 @@ def decode(token):
     return jwt.decode(token, settings.secret_key)
 
 
-def validate_frontend(token):
+def validate_frontend(token) -> Optional[FrontendToken]:
     try:
         data = decode(token)
-        assert data["type"] == "frontend"
-        return True
+        return FrontendToken.parse_obj(data)
 
-    except Exception as e:
+    except ValidationError as e:
         print(e)
-        return False
+        return None
 
 
-def validate_computer(token):
+def validate_computer(token) -> Optional[ComputerToken]:
     try:
         data = decode(token)
-        assert data["type"] == "computer"
-        return True
+        return ComputerToken.parse_obj(data)
 
-    except Exception as e:
+    except ValidationError as e:
         print(e)
-        return False
+        return None

backend/main.py

@@ -9,7 +9,7 @@ from pydantic import BaseModel
 from settings import settings
 import auth
 from state import StateManager
-from proto import Computer, ComputerType
+from proto import Computer, ComputerToken, ComputerType
 
 app = FastAPI()
 
@@ -18,7 +18,7 @@ state_manager = StateManager()
 
 @app.get("/api/{token}/validate", tags=["frontend"])
 async def validate_token(token: str):
-    return {"success": auth.validate_frontend(token)}
+    return {"success": auth.validate_frontend(token) is not None}
 
 
 @app.websocket("/api/{token}/state")
@@ -73,13 +73,32 @@ async def issue_new_token(data: RegistrationRequest):
     return {"token": auth.encode({"type": "computer", "uuid": str(computer.uuid)})}
 
 
+@app.websocket("/computer/{token}/socket")
+async def computer_websocket(websocket: WebSocket, token: str):
+
+    token = auth.validate_computer(token)
+
+    if not token:
+        await websocket.close()
+        return
+
+    await websocket.accept()
+    await state_manager.on_computer_connect(token.uuid)
+
+    try:
+        while True:
+            await websocket.receive_json()
+    except WebSocketDisconnect:
+        await state_manager.on_computer_disconnect(token.uuid)
+
+
 if settings.dev_mode:
 
     print("Starting in development mode.")
     print(f"Proxying requests to npm server on localhost:{settings.dev_npm_port}")
     import httpx
 
-    @app.get("/{path:path}")
+    @app.get("/{path:path}", tags=["dev mode please ignore"])
     async def dev_mode_proxy(path: str, response: Response):
 
         async with httpx.AsyncClient() as proxy:

backend/proto.py

@@ -5,10 +5,16 @@ from enum import Enum
 from pydantic import BaseModel
 
 
-class ComputerType(str, Enum):
-    COMPUTER = "computer"
-    TURTLE = "turtle"
-    POCKET = "pocket"
+class FrontendToken(BaseModel):
+    type: Literal["frontend"]
+
+
+class ComputerToken(BaseModel):
+    type: Literal["computer"]
+    uuid: UUID
+
+
+ComputerType = Literal["computer", "turtle", "pocket"]
 
 
 class Computer(BaseModel):
@@ -18,6 +24,21 @@ class Computer(BaseModel):
     label: Optional[str]
     group: str
 
+    # uUiD iS nOt JsOn SeRiAlIzAbLe
+    def dict(self):
+        data = super().dict()
+        data["uuid"] = str(data["uuid"])
+        return data
+
+
+class DynamicComputerState(BaseModel):
+    is_online: bool
+
+
+class StateItem(BaseModel):
+    static: Computer
+    dynamic: DynamicComputerState
+
 
 class State:
-    computers: List[Computer]
+    computers: List[StateItem]

backend/state.py

@@ -1,9 +1,10 @@
 import asyncio
+from uuid import UUID
 
 from tinydb import TinyDB
 
 from settings import settings
-from proto import Computer
+from proto import Computer, DynamicComputerState
 
 db = TinyDB(settings.database_path)
 computers = db.table("computers")
@@ -12,12 +13,26 @@ 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": computers.all()}
+
+        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:
@@ -36,12 +51,19 @@ class StateManager:
     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):
 
-        # uUiD iS nOt JsOn SeRiAlIzAbLe
-        computer_data = computer.dict()
-        computer_data["uuid"] = str(computer_data["uuid"])
-
-        computers.insert(computer_data)
+        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}")