From fcd8c4661a13d98dee292f9ee969909bd2bdf8d8 Mon Sep 17 00:00:00 2001
From: Pavel Veretennikov <vermut@pm.me>
Date: Wed, 8 May 2024 23:05:38 +0300
Subject: [PATCH] fix: tolerate unknown and invalid messages

---
 .../amc_alarm/amc_alarm_api/amc_proto.py      |  2 +-
 .../amc_alarm/amc_alarm_api/api.py            | 70 ++++++++++---------
 2 files changed, 39 insertions(+), 33 deletions(-)

diff --git a/custom_components/amc_alarm/amc_alarm_api/amc_proto.py b/custom_components/amc_alarm/amc_alarm_api/amc_proto.py
index 32e940a..e37e193 100644
--- a/custom_components/amc_alarm/amc_alarm_api/amc_proto.py
+++ b/custom_components/amc_alarm/amc_alarm_api/amc_proto.py
@@ -98,7 +98,7 @@ class AmcCommand(BaseModel):
 
 class AmcCommandResponse(BaseModel):
     command: str
-    status: str
+    status: Optional[str] = None
     centrals: Optional[dict[str, AmcCentralResponse]] = None
     user: Optional[AmcUser] = None
     token: Optional[str] = None
diff --git a/custom_components/amc_alarm/amc_alarm_api/api.py b/custom_components/amc_alarm/amc_alarm_api/api.py
index 0e5617c..b318c05 100644
--- a/custom_components/amc_alarm/amc_alarm_api/api.py
+++ b/custom_components/amc_alarm/amc_alarm_api/api.py
@@ -119,31 +119,7 @@ async def _running(self) -> None:
 
                         if message.type == aiohttp.WSMsgType.TEXT:
                             _LOGGER.debug("Websocket received data: %s", message.data)
-                            data = AmcCommandResponse.parse_raw(message.data)
-
-                            match data.command:
-                                case AmcCommands.LOGIN_USER:
-                                    if data.status == AmcCommands.STATUS_LOGGED_IN:
-                                        _LOGGER.debug("Authorized")
-                                        self._ws_state = ConnectionState.AUTHENTICATED
-                                        self._sessionToken = data.user.token
-                                    else:
-                                        _LOGGER.debug(
-                                            "Authorization failure: %s" % data.status
-                                        )
-                                        raise AuthenticationFailed(data.status)
-
-                                case AmcCommands.GET_STATES:
-                                    if data.status == AmcCommands.STATUS_OK:
-                                        self._raw_states = data.centrals
-                                    if self._callback:
-                                        await self._callback()
-                                    else:
-                                        _LOGGER.debug(
-                                            "Error getting _raw_states: %s"
-                                            % data.centrals
-                                        )
-                                        raise AmcException(data.centrals)
+                            await self._process_message(message)
 
             except aiohttp.ClientResponseError as error:
                 _LOGGER.error("Unexpected response received from server : %s", error)
@@ -172,13 +148,43 @@ async def _running(self) -> None:
                     _LOGGER.exception("Unexpected exception occurred: %s", error)
                     self._ws_state = ConnectionState.STOPPED
 
-    async def _login(self):
-        _LOGGER.info("Logging in with email: %s", self._login_email)
-        login_message = AmcCommand(
-            command=AmcCommands.LOGIN_USER,
-            data=AmcLogin(email=self._login_email, password=self._password),
-        )
-        await self._send_message(login_message)
+    async def _process_message(self, message):
+        try:
+            data = AmcCommandResponse.parse_raw(message.data)
+        except ValueError:
+            _LOGGER.warning("Can't process data from server: %s" % data)
+            return
+
+        match data.command:
+            case AmcCommands.LOGIN_USER:
+                if data.status == AmcCommands.STATUS_LOGGED_IN:
+                    _LOGGER.debug("Authorized")
+                    self._ws_state = ConnectionState.AUTHENTICATED
+                    self._sessionToken = data.user.token
+                    self._failed_attempts = 0
+                else:
+                    _LOGGER.debug("Authorization failure: %s" % data.status)
+                    raise AuthenticationFailed(data.status)
+
+            case AmcCommands.GET_STATES:
+                if data.status == AmcCommands.STATUS_OK:
+                    self._raw_states = data.centrals
+                if self._callback:
+                    await self._callback()
+                else:
+                    _LOGGER.debug("Error getting _raw_states: %s" % data.centrals)
+                    raise AmcException(data.centrals)
+            case _:
+                _LOGGER.warning("Unknown command received from server : %s", data)
+
+
+async def _login(self):
+    _LOGGER.info("Logging in with email: %s", self._login_email)
+    login_message = AmcCommand(
+        command=AmcCommands.LOGIN_USER,
+        data=AmcLogin(email=self._login_email, password=self._password),
+    )
+    await self._send_message(login_message)
 
     async def disconnect(self):
         _LOGGER.debug("Disconnecting")