Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: cleanup processing of websocket packets #24

Merged
merged 3 commits into from
Jun 11, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 24 additions & 33 deletions src/uiprotect/data/bootstrap.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,6 @@
from .user import Group, User
from .websocket import (
WSAction,
WSJSONPacketFrame,
WSPacket,
WSSubscriptionMessage,
)
Expand Down Expand Up @@ -169,6 +168,10 @@
id: str


_ModelType_NVR_value = ModelType.NVR.value
_ModelType_Event_value = ModelType.EVENT.value


class Bootstrap(ProtectBaseObject):
auth_user_id: str
access_key: str
Expand Down Expand Up @@ -536,65 +539,53 @@
models: set[ModelType] | None = None,
ignore_stats: bool = False,
) -> WSSubscriptionMessage | None:
if models is None:
models = set()

if not isinstance(packet.action_frame, WSJSONPacketFrame):
_LOGGER.debug(
"Unexpected action frame format: %s",
packet.action_frame.payload_format,
)

if not isinstance(packet.data_frame, WSJSONPacketFrame):
_LOGGER.debug(
"Unexpected data frame format: %s",
packet.data_frame.payload_format,
)
bdraco marked this conversation as resolved.
Show resolved Hide resolved

"""Process a WS packet."""
action, data = self._get_frame_data(packet)
if action["newUpdateId"] is not None:
self.last_update_id = action["newUpdateId"]
new_update_id: str = action["newUpdateId"]
if new_update_id is not None:
self.last_update_id = new_update_id

if action["modelKey"] not in ModelType.values_set():
_LOGGER.debug("Unknown model type: %s", action["modelKey"])
model_key: str = action["modelKey"]
if model_key not in ModelType.values_set():
_LOGGER.debug("Unknown model type: %s", model_key)

Check warning on line 550 in src/uiprotect/data/bootstrap.py

View check run for this annotation

Codecov / codecov/patch

src/uiprotect/data/bootstrap.py#L550

Added line #L550 was not covered by tests
self._create_stat(packet, None, True)
return None

if len(models) > 0 and ModelType(action["modelKey"]) not in models:
if models and ModelType(model_key) not in models:
self._create_stat(packet, None, True)
return None

if action["action"] == "remove":
action_action: str = action["action"]
if action_action == "remove":
return self._process_remove_packet(packet, data)

if data is None or len(data) == 0:
if not data:
self._create_stat(packet, None, True)
return None

try:
if action["action"] == "add":
if action_action == "add":
return self._process_add_packet(packet, data)

if action["action"] == "update":
if action["modelKey"] == ModelType.NVR.value:
if action_action == "update":
if model_key == _ModelType_NVR_value:
return self._process_nvr_update(packet, data, ignore_stats)

if (
action["modelKey"] in ModelType.bootstrap_models_set()
or action["modelKey"] == ModelType.EVENT.value
model_key in ModelType.bootstrap_models_set()
or model_key == _ModelType_Event_value
):
return self._process_device_update(
packet,
action,
data,
ignore_stats,
)
_LOGGER.debug(
"Unexpected bootstrap model type deviceadoptedfor update: %s",
action["modelKey"],
)
except (ValidationError, ValueError) as err:
self._handle_ws_error(action, err)

_LOGGER.debug(

Check warning on line 586 in src/uiprotect/data/bootstrap.py

View check run for this annotation

Codecov / codecov/patch

src/uiprotect/data/bootstrap.py#L586

Added line #L586 was not covered by tests
"Unexpected bootstrap model type deviceadoptedfor update: %s", model_key
bdraco marked this conversation as resolved.
Show resolved Hide resolved
)
bdraco marked this conversation as resolved.
Show resolved Hide resolved
self._create_stat(packet, None, True)
return None

Expand Down
12 changes: 9 additions & 3 deletions src/uiprotect/data/websocket.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import struct
import zlib
from dataclasses import dataclass
from functools import cache, cached_property
from typing import TYPE_CHECKING, Any

import orjson
Expand Down Expand Up @@ -62,6 +63,7 @@ def get_binary_from_data(self) -> bytes:
raise NotImplementedError

@staticmethod
@cache
def klass_from_format(format_raw: int) -> type[BaseWSPacketFrame]:
payload_format = ProtectWSPayloadFormat(format_raw)

Expand Down Expand Up @@ -174,13 +176,15 @@ def json(self) -> bytes:


class WSPacket:
"""Class to handle a unifi protect websocket packet."""

_raw: bytes
_raw_encoded: str | None = None

_action_frame: BaseWSPacketFrame | None = None
_data_frame: BaseWSPacketFrame | None = None

def __init__(self, data: bytes):
def __init__(self, data: bytes) -> None:
self._raw = data

def decode(self) -> None:
Expand All @@ -190,7 +194,7 @@ def decode(self) -> None:
self._action_frame.length,
)

@property
@cached_property
def action_frame(self) -> BaseWSPacketFrame:
if self._action_frame is None:
self.decode()
Expand All @@ -200,7 +204,7 @@ def action_frame(self) -> BaseWSPacketFrame:

return self._action_frame

@property
@cached_property
def data_frame(self) -> BaseWSPacketFrame:
if self._data_frame is None:
self.decode()
Expand All @@ -220,6 +224,8 @@ def raw(self, data: bytes) -> None:
self._action_frame = None
self._data_frame = None
self._raw_encoded = None
self.__dict__.pop("data_frame", None)
self.__dict__.pop("action_frame", None)

@property
def raw_base64(self) -> str:
Expand Down
Loading