websockets: add agent registration

Signed-off-by:  Eric Callahan <arksine.code@gmail.com>
This commit is contained in:
Eric Callahan 2022-04-08 08:27:52 -04:00
parent a79e0eeef2
commit cbc02d625c
No known key found for this signature in database
GPG Key ID: 7027245FBBDDF59A
1 changed files with 13 additions and 0 deletions

View File

@ -31,6 +31,7 @@ if TYPE_CHECKING:
from moonraker import Server from moonraker import Server
from app import APIDefinition from app import APIDefinition
from klippy_connection import KlippyConnection as Klippy from klippy_connection import KlippyConnection as Klippy
from .components.extensions import ExtensionManager
import components.authorization import components.authorization
_T = TypeVar("_T") _T = TypeVar("_T")
_C = TypeVar("_C", str, bool, float, int) _C = TypeVar("_C", str, bool, float, int)
@ -397,6 +398,14 @@ class WebsocketManager(APITransport):
"type": client_type, "type": client_type,
"url": url "url": url
} }
if client_type == "agent":
extensions: ExtensionManager
extensions = self.server.lookup_component("extensions")
try:
extensions.register_agent(ws)
except ServerError:
ws.client_data["type"] = ""
raise
logging.info( logging.info(
f"Websocket {ws.uid} Client Identified - " f"Websocket {ws.uid} Client Identified - "
f"Name: {name}, Version: {version}, Type: {client_type}" f"Name: {name}, Version: {version}, Type: {client_type}"
@ -629,6 +638,10 @@ class WebSocket(WebSocketHandler, Subscribable):
f"Close Code: {self.close_code}, " f"Close Code: {self.close_code}, "
f"Close Reason: {self.close_reason}, " f"Close Reason: {self.close_reason}, "
f"Pong Time Elapsed: {pong_elapsed:.2f}") f"Pong Time Elapsed: {pong_elapsed:.2f}")
if self._client_data["type"] == "agent":
extensions: ExtensionManager
extensions = self.server.lookup_component("extensions")
extensions.remove_agent(self)
self.wsm.remove_websocket(self) self.wsm.remove_websocket(self)
def check_origin(self, origin: str) -> bool: def check_origin(self, origin: str) -> bool: