Globally allow Aqara devices to handle swapped ZCL header direction (#4223)

* Globally allow Aqara devices to handle swapped ZCL header direction

* Use exceptions instead of verbose logic

* Add a test

Author: puddly

tests/test_xiaomi.py

@@ -2257,3 +2257,24 @@ def test_t1m_ceiling_light(zigpy_device_from_v2_quirk, endpoint):
         == AqaraLightT1M.AttributeDefs.power_on_state.id
     )
     assert cluster_listener.attribute_updates[1][1] == LumiPowerOnStateMode.Off
+
+
+async def test_lumi_magnet_sensor_aq2_bad_direction(zigpy_device_from_quirk):
+    """Test Aqara Magnet Sensor AQ2 quirk dealing with bad ZCL command direction."""
+
+    device = zigpy_device_from_quirk(zhaquirks.xiaomi.aqara.magnet_aq2.MagnetAQ2)
+    listener = ClusterListener(device.endpoints[1].out_clusters[OnOff.cluster_id])
+
+    # The device has a bad ZCL header and reports the incorrect direction for commands
+    device.packet_received(
+        t.ZigbeePacket(
+            profile_id=260,
+            cluster_id=6,
+            src_ep=1,
+            dst_ep=1,
+            data=t.SerializableBytes(bytes.fromhex("18930A00001001")),
+        )
+    )
+
+    # Our matching logic should be forgiving
+    assert listener.attribute_updates == [(0, t.Bool.true)]

zhaquirks/xiaomi/__init__.py

@@ -12,7 +12,7 @@
 from zigpy.profiles import zha
 from zigpy.quirks import CustomCluster, CustomDevice
 from zigpy.typing import AddressingMode
-from zigpy.zcl import foundation
+from zigpy.zcl import Cluster, foundation
 from zigpy.zcl.clusters.general import (
     AnalogInput,
     Basic,
@@ -122,6 +122,30 @@ def __init__(self, *args, **kwargs):
             self.battery_size = BatterySize.CR2032
         super().__init__(*args, **kwargs)
 
+    def _find_zcl_cluster(
+        self, hdr: foundation.ZCLHeader, packet: t.ZigbeePacket
+    ) -> Cluster:
+        """Find a cluster for the packet."""
+
+        # Aqara devices seem to be very lax with their ZCL header's `direction` field,
+        # we should try "flipping" it if matching doesn't work normally.
+        try:
+            return super()._find_zcl_cluster(hdr, packet)
+        except KeyError:
+            _LOGGER.debug(
+                "Packet is coming in the wrong direction for cluster %s on endpoint %s,"
+                " swapping direction and trying again"
+            )
+
+            return super()._find_zcl_cluster(
+                hdr.replace(
+                    frame_control=hdr.frame_control.replace(
+                        direction=hdr.frame_control.direction.flip()
+                    )
+                ),
+                packet,
+            )
+
 
 class XiaomiQuickInitDevice(XiaomiCustomDevice, QuickInitDevice):
     """Xiaomi devices eligible for QuickInit."""

zhaquirks/xiaomi/aqara/magnet_agl02.py

@@ -1,8 +1,6 @@
 """Quirk for Aqara T1 door sensor lumi.magnet.agl02."""
 
 from zigpy.profiles import zha
-import zigpy.types as t
-from zigpy.zcl import Cluster, foundation
 from zigpy.zcl.clusters.general import Basic, Identify, OnOff, Ota, PowerConfiguration
 from zigpy.zcl.clusters.security import IasZone
 
@@ -31,20 +29,6 @@ def __init__(self, *args, **kwargs):
         self.battery_size = BatterySize.CR1632
         super().__init__(*args, **kwargs)
 
-    def _find_zcl_cluster(
-        self, hdr: foundation.ZCLHeader, packet: t.ZigbeePacket
-    ) -> Cluster:
-        """Find a cluster for the packet."""
-        assert packet.src_ep is not None
-        endpoint = self.endpoints[packet.src_ep]
-
-        # The `direction` field is incorrectly set, we should always route the packet
-        # to client `OnOff` cluster
-        if packet.cluster_id == OnOff.cluster_id:
-            return endpoint.out_clusters[packet.cluster_id]
-
-        return super()._find_zcl_cluster(hdr, packet)
-
     signature = {
         MODELS_INFO: [("LUMI", "lumi.magnet.agl02")],
         ENDPOINTS: {