Add data drop when needed (#357)

Vaskivskyi · Nov 9, 2023 · 42ed7f3 · 42ed7f3
1 parent d7c705b
commit 42ed7f3
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 10 deletions.
diff --git a/asusrouter/asusrouter.py b/asusrouter/asusrouter.py
@@ -428,6 +428,19 @@ def _transform_data(self, datatype: AsusData, data: Any) -> Any:
 
         return data
 
+    def _drop_data(self, datatype: AsusData, endpoint: Endpoint) -> bool:
+        """Check whether data should be dropped. This is required
+        for some data obtained from multiple endpoints."""
+
+        if not self._identity:
+            return False
+
+        if datatype == AsusData.OPENVPN_CLIENT:
+            if self._identity.merlin is True:
+                return endpoint == Endpoint.HOOK
+
+        return False
+
     def _check_state(self, datatype: Optional[AsusData]) -> None:
         """Make sure the state object is available."""
 
@@ -513,17 +526,24 @@ async def async_get_data(self, datatype: AsusData, force: bool = False) -> Any:
                 if not self._identity:
                     self._identity = await self.async_get_identity()
 
-                result = merge_dicts(
-                    result,
-                    process(
-                        endpoint,
-                        data,
-                        self._state,
-                        self._identity.firmware,
-                        self._identity.wlan,
-                    ),
+                processed = process(
+                    endpoint,
+                    data,
+                    self._state,
+                    self._identity.firmware,
+                    self._identity.wlan,
                 )
 
+                # Check whether data should be dropped
+                to_drop = []
+                for key, value in processed.items():
+                    if self._drop_data(key, endpoint):
+                        to_drop.append(key)
+                for key in to_drop:
+                    processed.pop(key, None)
+
+                result = merge_dicts(result, processed)
+
                 # Check if we have data and data finder merge is ANY
                 if result and data_finder.merge == AsusDataMerge.ANY:
                     break

diff --git a/asusrouter/modules/endpoint/hook_const.py b/asusrouter/modules/endpoint/hook_const.py
@@ -4,7 +4,7 @@
 
 from asusrouter.modules.ip_address import read_dns_ip_address, read_ip_address_type
 from asusrouter.modules.openvpn import AsusOVPNServer
-from asusrouter.modules.wireguard import AsusWireGuardServer
+from asusrouter.modules.wireguard import AsusWireGuardClient, AsusWireGuardServer
 from asusrouter.tools.converters import (
     safe_bool,
     safe_int,