diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py index 068034a6..888ab562 100644 --- a/bellows/zigbee/application.py +++ b/bellows/zigbee/application.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from datetime import datetime, timezone import logging import os import statistics @@ -93,6 +94,7 @@ def __init__(self, config: dict): self._watchdog_feed_counter = 0 self._req_lock = asyncio.Lock() + self._packet_capture_channel: int | None = None @property def controller_event(self): @@ -752,6 +754,48 @@ async def _network_scan( rssi=lastHopRssi, ) + def _check_status(self, status: t.sl_Status | t.EmberStatus) -> None: + if t.sl_Status.from_ember_status(status) != t.sl_Status.OK: + raise ControllerError(f"Command failed: {status!r}") + + async def _packet_capture(self, channel: int): + (status,) = await self._ezsp.mfglibStart(rxCallback=True) + self._check_status(status) + + try: + await self._packet_capture_change_channel(channel=channel) + assert self._packet_capture_channel is not None + + queue = asyncio.Queue() + + with self._ezsp.callback_for_commands( + {"mfglibRxHandler"}, + callback=lambda _, response: queue.put_nowait( + (datetime.now(timezone.utc), response) + ), + ): + while True: + timestamp, (linkQuality, rssi, packetContents) = await queue.get() + + # The last two bytes are not a FCS + packetContents = packetContents[:-2] + + yield zigpy.types.CapturedPacket( + timestamp=timestamp, + rssi=rssi, + lqi=linkQuality, + channel=self._packet_capture_channel, + data=packetContents, + ) + finally: + (status,) = await self._ezsp.mfglibEnd() + self._check_status(status) + + async def _packet_capture_change_channel(self, channel: int): + (status,) = await self._ezsp.mfglibSetChannel(channel=channel) + self._check_status(status) + self._packet_capture_channel = channel + async def send_packet(self, packet: zigpy.types.ZigbeePacket) -> None: if not self.is_controller_running: raise ControllerError("ApplicationController is not running") diff --git a/pyproject.toml b/pyproject.toml index 3a9c9b1a..6256ca23 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,7 @@ dependencies = [ "click-log>=0.2.1", "pure_pcapy3==1.0.1", "voluptuous", - "zigpy>=0.70.0", + "zigpy>=0.75.0", 'async-timeout; python_version<"3.11"', ] diff --git a/tests/test_application.py b/tests/test_application.py index 873f6921..c0baa95d 100644 --- a/tests/test_application.py +++ b/tests/test_application.py @@ -2044,3 +2044,75 @@ async def test_network_scan_failure(app: ControllerApplication) -> None: channels=t.Channels.from_channel_list([11, 15, 26]), duration_exp=4 ): pass + + +async def test_packet_capture(app: ControllerApplication) -> None: + app._ezsp._protocol.mfglibStart.return_value = [t.sl_Status.OK] + app._ezsp._protocol.mfglibSetChannel.return_value = [t.sl_Status.OK] + app._ezsp._protocol.mfglibEnd.return_value = [t.sl_Status.OK] + + async def receive_packets() -> None: + app._ezsp._protocol._handle_callback( + "mfglibRxHandler", + list( + { + "linkQuality": 150, + "rssi": -70, + "packetContents": b"packet 1\xAB\xCD", + }.values() + ), + ) + + await asyncio.sleep(0.5) + + app._ezsp._protocol._handle_callback( + "mfglibRxHandler", + list( + { + "linkQuality": 200, + "rssi": -50, + "packetContents": b"packet 2\xAB\xCD", + }.values() + ), + ) + + task = asyncio.create_task(receive_packets()) + packets = [] + + async for packet in app.packet_capture(channel=15): + packets.append(packet) + + if len(packets) == 1: + await app.packet_capture_change_channel(channel=20) + elif len(packets) == 2: + break + + assert packets == [ + zigpy_t.CapturedPacket( + timestamp=packets[0].timestamp, + rssi=-70, + lqi=150, + channel=15, + data=b"packet 1", + ), + zigpy_t.CapturedPacket( + timestamp=packets[1].timestamp, + rssi=-50, + lqi=200, + channel=20, # The second packet's channel was changed + data=b"packet 2", + ), + ] + + await task + await asyncio.sleep(0.1) + + assert app._ezsp._protocol.mfglibEnd.mock_calls == [call()] + + +async def test_packet_capture_failure(app: ControllerApplication) -> None: + app._ezsp._protocol.mfglibStart.return_value = [t.sl_Status.FAIL] + + with pytest.raises(zigpy.exceptions.ControllerException): + async for packet in app.packet_capture(channel=15): + pass