From d97058f31e170b75c7f45be1b1bd25f0a6f8e521 Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Mon, 19 Aug 2024 15:33:46 -0400
Subject: [PATCH] Implement `network_scan`

---
 bellows/zigbee/application.py | 38 +++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py
index 22247178..35834f91 100644
--- a/bellows/zigbee/application.py
+++ b/bellows/zigbee/application.py
@@ -5,6 +5,7 @@
 import os
 import statistics
 import sys
+from typing import AsyncGenerator
 
 if sys.version_info[:2] < (3, 11):
     from async_timeout import timeout as asyncio_timeout  # pragma: no cover
@@ -698,6 +699,43 @@ async def energy_scan(
             for channel in list(channels)
         }
 
+    async def network_scan(
+        self, channels: t.Channels, duration: int
+    ) -> AsyncGenerator[zigpy.types.NetworkBeacon]:
+        """Scans for networks and yields network beacons."""
+        queue = asyncio.Queue()
+
+        with self._ezsp.callback_for_commands(
+            {"networkFoundHandler", "scanCompleteHandler"},
+            callback=lambda command, response: queue.put_nowait((command, response)),
+        ):
+            # XXX: replace with normal command invocation once overload is removed
+            (status,) = await self._ezsp._command(
+                "startScan",
+                scanType=t.EzspNetworkScanType.ACTIVE_SCAN,
+                channelMask=channels,
+                duration=duration,
+            )
+
+            while True:
+                command, response = await queue.get()
+
+                if command == "scanCompleteHandler":
+                    break
+
+                (networkFound, lastHopLqi, lastHopRssi) = response
+
+                yield zigpy.types.NetworkBeacon(
+                    pan_id=networkFound.panId,
+                    extended_pan_id=networkFound.extendedPanId,
+                    channel=networkFound.channel,
+                    nwk_update_id=networkFound.nwkUpdateId,
+                    permit_joining=bool(networkFound.allowingJoin),
+                    stack_profile=networkFound.stackProfile,
+                    lqi=lastHopLqi,
+                    rssi=lastHopRssi,
+                )
+
     async def send_packet(self, packet: zigpy.types.ZigbeePacket) -> None:
         if not self.is_controller_running:
             raise ControllerError("ApplicationController is not running")