Refactor device join. (#31)

Adminiuga · web-flow · commit 452fc19d431f · 2019-03-12T20:07:49.000-04:00
Use ZDO device_annce message to trigger new device join.
diff --git a/tests/test_application.py b/tests/test_application.py
@@ -3,7 +3,7 @@
 
 import pytest
 
-from zigpy.types import EUI64
+from zigpy.types import EUI64, uint16_t
 from zigpy_xbee.api import ModemStatus, XBee
 from zigpy_xbee.zigbee.application import ControllerApplication
 
@@ -20,19 +20,22 @@ def test_modem_status(app):
     app.handle_modem_status(ModemStatus(0xEE))
 
 
-def _test_rx(app, device, deserialized):
+def _test_rx(app, device, nwk, deserialized,
+             dst_ep=mock.sentinel.dst_ep,
+             cluster_id=mock.sentinel.cluster_id,
+             data=b''):
     app.get_device = mock.MagicMock(return_value=device)
     app.deserialize = mock.MagicMock(return_value=deserialized)
 
     app.handle_rx(
         b'\x01\x02\x03\x04\x05\x06\x07\x08',
-        mock.sentinel.src_nwk,
+        nwk,
         mock.sentinel.src_ep,
-        mock.sentinel.dst_ep,
-        mock.sentinel.cluster_id,
+        dst_ep,
+        cluster_id,
         mock.sentinel.profile_id,
         mock.sentinel.rxopts,
-        b'',
+        data,
     )
 
     assert app.deserialize.call_count == 1
@@ -41,7 +44,7 @@ def _test_rx(app, device, deserialized):
 def test_rx(app):
     device = mock.MagicMock()
     app.handle_message = mock.MagicMock()
-    _test_rx(app, device, (1, 2, False, []))
+    _test_rx(app, device, mock.sentinel.src_nwk, (1, 2, False, []))
     assert app.handle_message.call_count == 1
     assert app.handle_message.call_args == ((
         device,
@@ -73,16 +76,43 @@ def test_rx_nwk_0000(app):
     assert app._handle_reply.call_count == 0
 
 
+def test_rx_unknown_device(app):
+    app._handle_reply = mock.MagicMock()
+    app.handle_message = mock.MagicMock()
+    app.handle_join = mock.MagicMock()
+    app.get_device = mock.MagicMock(side_effect=KeyError)
+    app.handle_rx(
+        b'\x01\x02\x03\x04\x05\x06\x07\x08',
+        0x1234,
+        mock.sentinel.src_ep,
+        mock.sentinel.dst_ep,
+        mock.sentinel.cluster_id,
+        mock.sentinel.profile_id,
+        mock.sentinel.rxopts,
+        b''
+    )
+    assert app.handle_join.call_count == 0
+    assert app.get_device.call_count == 1
+    assert app.handle_message.call_count == 0
+    assert app._handle_reply.call_count == 0
+
+
 def test_rx_reply(app):
     app._handle_reply = mock.MagicMock()
-    _test_rx(app, mock.MagicMock(), (1, 2, True, []))
+    _test_rx(app, mock.MagicMock(), mock.sentinel.src_nwk, (1, 2, True, []))
     assert app._handle_reply.call_count == 1
 
 
 def test_rx_failed_deserialize(app, caplog):
+    from zigpy.device import Status as DeviceStatus
+
     app._handle_reply = mock.MagicMock()
     app.handle_message = mock.MagicMock()
-    app.get_device = mock.MagicMock(return_value=mock.sentinel.device)
+    nwk = 0x1234
+    ieee, _ = EUI64.deserialize(b'\x01\x02\x03\x04\x05\x06\x07\x08')
+    device = app.add_device(ieee, nwk)
+    device.status = DeviceStatus.ENDPOINTS_INIT
+    app.get_device = mock.MagicMock(return_value=device)
     app.deserialize = mock.MagicMock(side_effect=ValueError)
 
     app.handle_rx(
@@ -102,6 +132,91 @@ def test_rx_failed_deserialize(app, caplog):
     assert app.handle_message.call_count == 0
 
 
+@pytest.fixture
+def device(app):
+    def _device(new=False, zdo_init=False, nwk=uint16_t(0x1234)):
+        from zigpy.device import Device, Status as DeviceStatus
+
+        ieee, _ = EUI64.deserialize(b'\x08\x07\x06\x05\x04\x03\x02\x01')
+        dev = Device(app, ieee, nwk)
+        if new:
+            dev.status = DeviceStatus.NEW
+        elif zdo_init:
+            dev.status = DeviceStatus.ZDO_INIT
+        else:
+            dev.status = DeviceStatus.ENDPOINTS_INIT
+        return dev
+    return _device
+
+
+def _device_join(app, dev, data):
+    app.handle_message = mock.MagicMock()
+    app.handle_join = mock.MagicMock()
+
+    deserialized = (1, 2, False, [])
+    dst_ep = 0
+    cluster_id = 0x0013
+
+    _test_rx(app, dev, dev.nwk, deserialized, dst_ep, cluster_id, data)
+    assert app.handle_join.call_count == 1
+    assert app.handle_message.call_count == 1
+
+
+def test_device_join_new(app, device):
+    dev = device()
+    data = b'\xee' + dev.nwk.serialize() + dev.ieee.serialize()
+
+    _device_join(app, dev, data)
+
+
+def test_device_join_inconsistent_nwk(app, device):
+    dev = device()
+    data = b'\xee' + b'\x01\x02' + dev.ieee.serialize()
+
+    _device_join(app, dev, data)
+
+
+def test_device_join_inconsistent_ieee(app, device):
+    dev = device()
+    data = b'\xee' + dev.nwk.serialize() + b'\x01\x02\x03\x04\x05\x06\x07\x08'
+
+    _device_join(app, dev, data)
+
+
+def _device_status_new(app, dev):
+    app.handle_join = mock.MagicMock()
+    app.handle_message = mock.MagicMock()
+    app.handle_reply = mock.MagicMock()
+    app.get_device = mock.MagicMock(return_value=dev)
+    app.deserialize = mock.MagicMock()
+
+    app.handle_rx(
+        b'\x01\x02\x03\x04\x05\x06\x07\x08',
+        dev.nwk,
+        mock.sentinel.src_ep,
+        mock.sentinel.dst_ep,
+        mock.sentinel.cluster_id,
+        mock.sentinel.profile_id,
+        mock.sentinel.rxopts,
+        b'',
+    )
+
+    assert app.deserialize.call_count == 0
+    assert app.handle_join.call_count == 0
+    assert app.handle_message.call_count == 0
+    assert app.handle_reply.call_count == 0
+
+
+def test_handle_rx_device_status_new(app, device):
+    dev = device(new=True)
+    _device_status_new(app, dev)
+
+
+def test_handle_rx_device_status_zdo(app, device):
+    dev = device(zdo_init=True)
+    _device_status_new(app, dev)
+
+
 @pytest.mark.asyncio
 async def test_broadcast(app):
     (profile, cluster, src_ep, dst_ep, grpid, radius, tsn, data) = (
@@ -295,7 +410,8 @@ async def test_permit(app):
 async def _test_request(app, do_reply=True, expect_reply=True, **kwargs):
     seq = 123
     nwk = 0x2345
-    app._devices_by_nwk[nwk] = 0x22334455
+    ieee = EUI64(b'\x01\x02\x03\x04\x05\x06\x07\x08')
+    app.add_device(ieee, nwk)
 
     def _mock_seq_command(cmdname, ieee, nwk, src_ep, dst_ep, cluster,
                           profile, radius, options, data):
diff --git a/zigpy_xbee/zigbee/application.py b/zigpy_xbee/zigbee/application.py
@@ -22,8 +22,6 @@ def __init__(self, api, database_file=None):
         api.set_application(self)
 
         self._pending = {}
-        self._devices_by_nwk = {d.nwk: d.ieee for a, d in self.devices.items()}
-
         self._nwk = 0
 
     async def startup(self, auto_form=False):
@@ -116,9 +114,11 @@ async def request(self, nwk, profile, cluster, src_ep, dst_ep, sequence, data, e
         if expect_reply:
             reply_fut = asyncio.Future()
             self._pending[sequence] = reply_fut
+
+        dev = self.get_device(nwk=nwk)
         self._api._seq_command(
             'tx_explicit',
-            self._devices_by_nwk[nwk],
+            dev.ieee,
             nwk,
             src_ep,
             dst_ep,
@@ -153,10 +153,37 @@ def handle_rx(self, src_ieee, src_nwk, src_ep, dst_ep, cluster_id, profile_id, r
             return
 
         ember_ieee = zigpy.types.EUI64(src_ieee)
-        if ember_ieee not in self.devices:
-            self.handle_join(src_nwk, ember_ieee, 0)  # TODO: Parent nwk
-        self._devices_by_nwk[src_nwk] = src_ieee
-        device = self.get_device(ember_ieee)
+        if dst_ep == 0 and cluster_id == 0x13:
+            # ZDO Device announce request
+            nwk, data = zigpy.types.uint16_t.deserialize(data[1:])
+            ieee, data = zigpy.types.EUI64.deserialize(data)
+            LOGGER.info("New device joined: NWK 0x%04x, IEEE %s", nwk, ieee)
+            if ember_ieee != ieee:
+                LOGGER.warning(
+                    "Announced IEEE %s is different from originator %s",
+                    str(ieee), str(ember_ieee))
+            if src_nwk != nwk:
+                LOGGER.warning(
+                    "Announced 0x%04x NWK is different from originator 0x%04x",
+                    nwk, src_nwk
+                )
+            self.handle_join(nwk, ieee, 0)
+
+        try:
+            device = self.get_device(ieee=ember_ieee)
+        except KeyError:
+            LOGGER.debug("Received frame from unknown device: 0x%04x/%s",
+                         src_nwk, str(ember_ieee))
+            return
+
+        if device.status == zigpy.device.Status.NEW and dst_ep != 0:
+            # only allow ZDO responses while initializing device
+            LOGGER.debug("Received frame on uninitialized device %s (%s) for endpoint: %s", device.ieee, device.status, dst_ep)
+            return
+        elif device.status == zigpy.device.Status.ZDO_INIT and dst_ep != 0 and cluster_id != 0:
+            # only allow access to basic cluster while initializing endpoints
+            LOGGER.debug("Received frame on uninitialized device %s endpoint %s for cluster: %s", device.ieee, dst_ep, cluster_id)
+            return
 
         try:
             tsn, command_id, is_reply, args = self.deserialize(device, src_ep, cluster_id, data)
