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 FrontendToken.parse_obj(data)
-        return True
 
-    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 ComputerToken.parse_obj(data)
-        return True
 
-    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):
+class FrontendToken(BaseModel):
-    COMPUTER = "computer"
+    type: Literal["frontend"]
-    TURTLE = "turtle"
+
-    POCKET = "pocket"
+
+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
+        computers.insert(computer.dict())
-        computer_data = computer.dict()
-        computer_data["uuid"] = str(computer_data["uuid"])
-
-        computers.insert(computer_data)
         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}")

frontend/src/pages/Index.tsx

@@ -2,7 +2,7 @@ import { IconProp } from "@fortawesome/fontawesome-svg-core";
 import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import { fas } from "fontawesome.macro";
 import React, { useEffect, useState } from "react";
-import { Computer, ComputerType, State } from "../proto";
+import { Computer, ComputerType, State, StateItem } from "../proto";
 import { TokenContext } from "../tokenStorage";
 
 export const Index: React.FC = () => {
@@ -37,14 +37,14 @@ export const Index: React.FC = () => {
 
 export default Index;
 
-const Groups: React.FC<{ computers: Array<Computer> }> = ({ computers }) => {
+const Groups: React.FC<{ computers: Array<StateItem> }> = ({ computers }) => {
-    const groupMap = new Map<string, Array<Computer>>();
+    const groupMap = new Map<string, Array<StateItem>>();
 
     for (const computer of computers) {
-        if (!groupMap.has(computer.group)) {
+        if (!groupMap.has(computer.static.group)) {
-            groupMap.set(computer.group, []);
+            groupMap.set(computer.static.group, []);
         }
-        groupMap.get(computer.group)!.push(computer);
+        groupMap.get(computer.static.group)!.push(computer);
     }
 
     return <div>
@@ -61,15 +61,15 @@ const Groups: React.FC<{ computers: Array<Computer> }> = ({ computers }) => {
 }
 
 
-const CardList: React.FC<{ computers: Array<Computer> }> = ({ computers }) => {
+const CardList: React.FC<{ computers: Array<StateItem> }> = ({ computers }) => {
     return <>
         {computers.map(
-            (computer) => <ComputerCard key={computer.uuid} computer={computer} />
+            (computer) => <ComputerCard key={computer.static.uuid} computer={computer} />
         )}
     </>
 }
 
-const ComputerCard: React.FC<{ computer: Computer }> = ({ computer }) => {
+const ComputerCard: React.FC<{ computer: StateItem }> = ({ computer }) => {
 
     const typeIcon: Map<ComputerType, IconProp> = new Map([
         ["computer", fas`desktop`],
@@ -82,21 +82,33 @@ const ComputerCard: React.FC<{ computer: Computer }> = ({ computer }) => {
             <div>
                 <p>
                     <span className="inline-block align-middle">
-                        <FontAwesomeIcon icon={typeIcon.get(computer.type) || fas`question`} />
+                        <FontAwesomeIcon icon={typeIcon.get(computer.static.type) || fas`question`} />
                     </span>
-                    <span className={`heading mx-2 text-lg ${computer.label ? "" : "font-extralight"}`}>
+                    <span className={`heading mx-2 text-lg ${computer.static.label ? "" : "font-extralight"}`}>
-                        {computer.label || "(no label)"}
+                        {computer.static.label || "(no label)"}
                     </span>
-                    {computer.is_advanced && <span className="badge ~yellow @low">advanced</span>}
+                    {computer.static.is_advanced && <span className="badge ~yellow @low">advanced</span>}
                 </p>
                 <p className="support">
-                    {computer.uuid}
+                    {computer.static.uuid}
                 </p>
             </div>
             <div>
-                Offline
+                {
-                <span className="icon m-1 text-red-900"><FontAwesomeIcon icon={fas`ban`} /></span>
+                    computer.dynamic.is_online
+                        ? <>
+                            Online
+                            <span className="icon m-1 text-green-600"><FontAwesomeIcon icon={fas`circle`} /></span>
+                        </>
+                        : <>
+                            Offline
+                            <span className="icon m-1 text-red-900"><FontAwesomeIcon icon={fas`ban`} /></span>
+                        </>
+                }
             </div>
         </div>
-    </div>
+        <pre className="pre overflow-x-scroll">
+            {JSON.stringify(computer, null, 2)}
+        </pre>
+    </div >
 }

frontend/src/proto.ts

@@ -11,6 +11,15 @@ export type Computer = {
     group: string,
 }
 
+export type DynamicComputerState = {
+    is_online: boolean,
+}
+
+export type StateItem = {
+    static: Computer,
+    dynamic: DynamicComputerState,
+}
+
 export type State = {
-    computers: Array<Computer>
+    computers: Array<StateItem>
 }