From 740827678eb5d85a8b51e0cac7fd2e6610de1dd8 Mon Sep 17 00:00:00 2001
From: DamKas <48238600+DamKast@users.noreply.github.com>
Date: Wed, 13 Mar 2024 18:53:25 +0400
Subject: [PATCH] avoid _unknown_addr_discovery rewrite

---
 zigpy_zboss/zigbee/application.py | 90 ++++++++++++++++++++++++-------
 1 file changed, 70 insertions(+), 20 deletions(-)

diff --git a/zigpy_zboss/zigbee/application.py b/zigpy_zboss/zigbee/application.py
index 49ce664..5694769 100644
--- a/zigpy_zboss/zigbee/application.py
+++ b/zigpy_zboss/zigbee/application.py
@@ -504,26 +504,6 @@ def add_device(self, ieee: t.EUI64, nwk: t.NWK):
         self.devices[ieee] = dev
         return dev
 
-    async def _discover_unknown_device(self, nwk: t.NWK) -> None:
-        """Discover the IEEE address of a device with an unknown NWK."""
-        res = await self._api.request(
-            c.ZDO.IeeeAddrReq.Req(
-                TSN=self.get_sequence(),
-                DstNWK=t.BroadcastAddress.RX_ON_WHEN_IDLE,
-                NWKtoMatch=nwk,
-                RequestType=zdo_t.AddrRequestType.Single,
-                StartIndex=0,
-                )
-        )
-
-        nwk = res.RemoteDevNWK
-        ieee = res.RemoteDevIEEE
-
-        LOGGER.debug("Discovered IEEE address for NWK=%s: %s", nwk, ieee)
-        self.handle_join(
-            nwk=nwk, ieee=ieee, parent_nwk=None, handle_rejoin=False
-        )
-
     #####################################################
     # Callbacks attached during startup                 #
     #####################################################
@@ -645,6 +625,10 @@ async def send_packet(self, packet: t.ZigbeePacket) -> None:
             raise DeliveryError(
                 "Coordinator is disconnected, cannot send request")
 
+        if zigpy.zdo.ZDO_ENDPOINT in (packet.src_ep, packet.dst_ep):
+            await self._handle_zdo_packet(packet)
+            return
+
         LOGGER.debug("Sending packet %r", packet)
 
         options = c.aps.TransmitOptions.NONE
@@ -697,3 +681,69 @@ async def send_packet(self, packet: t.ZigbeePacket) -> None:
                     Payload=t_zboss.Payload(packet.data.serialize()),
                 )
             )
+
+    async def _handle_zdo_packet(self, packet: t.ZigbeePacket) -> None:
+        """ZDO packets that can't be send using the ZBOSS APSDE request."""
+        # The current zigpy device may not exist if we receive a packet early
+        try:
+            zdo = self._device.zdo
+        except KeyError:
+            zdo = zigpy.zdo.ZDO(None)
+
+        try:
+            zdo_hdr, zdo_args = zdo.deserialize(
+                cluster_id=packet.cluster_id, data=packet.data.serialize()
+            )
+        except ValueError:
+            LOGGER.debug("Could not parse ZDO message from packet")
+            return
+
+        if zdo_hdr.command_id == zdo_t.ZDOCmd.IEEE_addr_req:
+            await self._ieee_addr_req(packet, zdo_hdr, zdo_args)
+
+    async def _ieee_addr_req(
+            self,
+            packet: t.ZigbeePacket,
+            zdo_hdr: zdo_t.ZDOHeader,
+            zdo_args: tuple[Any]) -> None:
+        """Send ZDO IEEE addr request and handle the response."""
+        tsn = zdo_hdr.tsn
+        nwki, req_type, index = zdo_args
+        res = await self._api.request(
+            c.ZDO.IeeeAddrReq.Req(
+                TSN=tsn,
+                DstNWK=packet.dst.address,
+                NWKtoMatch=nwki,
+                RequestType=req_type,
+                StartIndex=index,
+                )
+        )
+
+        status = zdo_t.Status(res.StatusCode)
+        ieee = res.RemoteDevIEEE
+        nwki = res.RemoteDevNWK
+        data = tsn.serialize() \
+            + status.serialize() \
+            + ieee.serialize() \
+            + nwki.serialize()
+
+        packet = t.ZigbeePacket(
+            src=t.AddrModeAddress(
+                addr_mode=t.AddrMode.NWK,
+                address=res.RemoteDevNWK,
+            ),
+            src_ep=0,
+            dst=t.AddrModeAddress(
+                addr_mode=t.AddrMode.NWK,
+                address=self.state.node_info.nwk,
+            ),
+            dst_ep=0,
+            tsn=tsn,
+            profile_id=0,
+            cluster_id=zdo_t.ZDOCmd.IEEE_addr_rsp,
+            data=t.SerializableBytes(data),
+            tx_options=t.TransmitOptions.NONE,
+            lqi=None,
+            rssi=None
+        )
+        self.packet_received(packet)