Allow searching for bridges on more than one network

aiocomfoconnect/__main__.py

@@ -28,7 +28,7 @@
 async def main(args):
     """Main function."""
     if args.action == "discover":
-        await run_discover(args.host)
+        await run_discover(args.host, args.broadcast)
 
     elif args.action == "register":
         await run_register(args.host, args.uuid, args.name, args.pin)
@@ -67,9 +67,9 @@ async def main(args):
         raise UnknownActionException("Unknown action: " + args.action)
 
 
-async def run_discover(host: str = None):
+async def run_discover(host: str = None, broadcast_addresses: list = None):
     """Discover all bridges on the network."""
-    bridges = await discover_bridges(host)
+    bridges = await discover_bridges(host, broadcast_addresses=broadcast_addresses)
     print("Discovered bridges:")
     for bridge in bridges:
         print(bridge)
@@ -389,6 +389,7 @@ def main_cli():
 
     p_discover = subparsers.add_parser("discover", help="discover ComfoConnect LAN C devices on your network")
     p_discover.add_argument("--host", help="Host address of the bridge")
+    p_discover.add_argument("--broadcast", help="Broadcast address to search (eg. 192.168.1.255), can be repeated", action="append")
 
     p_register = subparsers.add_parser("register", help="register on a ComfoConnect LAN C device")
     p_register.add_argument("--pin", help="PIN code to register on the bridge", default=DEFAULT_PIN)

aiocomfoconnect/discovery.py

@@ -4,40 +4,54 @@
 
 import asyncio
 import logging
-from typing import Any, List, Union
+from typing import Any, Iterable, List, Union
 
 from .bridge import Bridge
 from .protobuf import zehnder_pb2
 
 _LOGGER = logging.getLogger(__name__)
 
+# The limited broadcast address. The operating system sends this out over a single interface, the one
+# that the routing table selects, so it doesn't reach bridges that live behind another interface.
+BROADCAST_ADDRESS = "<broadcast>"
+
 
 class BridgeDiscoveryProtocol(asyncio.DatagramProtocol):
     """UDP Protocol for the ComfoConnect LAN C bridge discovery."""
 
-    def __init__(self, target: str = None, timeout: int = 5):
+    def __init__(self, target: str = None, timeout: int = 5, broadcast_addresses: Iterable[Any] = None):
         loop = asyncio.get_running_loop()
 
         self._bridges: List[Bridge] = []
+        self._bridge_uuids: set[str] = set()
         self._target = target
         self._future = loop.create_future()
         self.transport = None
         self._timeout = loop.call_later(timeout, self.disconnect)
 
+        if target:
+            self._targets = [target]
+        elif broadcast_addresses:
+            # Addresses can be passed as ipaddress objects, since that is what Home Assistant hands out.
+            self._targets = [str(address) for address in broadcast_addresses]
+        else:
+            self._targets = [BROADCAST_ADDRESS]
+
     def connection_made(self, transport: asyncio.transports.DatagramTransport):
         """Called when a connection is made."""
         _LOGGER.debug("Socket has been created")
         self.transport = transport
 
-        if self._target:
-            _LOGGER.debug("Sending discovery request to %s:%d", self._target, Bridge.PORT)
-            self.transport.sendto(b"\x0a\x00", (self._target, Bridge.PORT))
-        else:
-            _LOGGER.debug("Sending discovery request to broadcast:%d", Bridge.PORT)
-            self.transport.sendto(b"\x0a\x00", ("<broadcast>", Bridge.PORT))
+        for target in self._targets:
+            _LOGGER.debug("Sending discovery request to %s:%d", target, Bridge.PORT)
+            self.transport.sendto(b"\x0a\x00", (target, Bridge.PORT))
 
     def datagram_received(self, data: Union[bytes, str], addr: tuple[str | Any, int]):
         """Called when some datagram is received."""
+        if self._future.done():
+            _LOGGER.debug("Ignoring data received from %s after the discovery finished", addr)
+            return
+
         if data == b"\x0a\x00":
             _LOGGER.debug("Ignoring discovery request from %s:%d", addr[0], addr[1])
             return
@@ -48,13 +62,18 @@ def datagram_received(self, data: Union[bytes, str], addr: tuple[str | Any, int]
         parser = zehnder_pb2.DiscoveryOperation()  # pylint: disable=no-member
         parser.ParseFromString(data)
 
-        self._bridges.append(
-            Bridge(
-                host=parser.searchGatewayResponse.ipaddress,
-                uuid=parser.searchGatewayResponse.uuid.hex(),
-                bridge_type=parser.searchGatewayResponse.type,
+        uuid = parser.searchGatewayResponse.uuid.hex()
+
+        # A bridge can reply to more than one of our discovery requests.
+        if uuid not in self._bridge_uuids:
+            self._bridge_uuids.add(uuid)
+            self._bridges.append(
+                Bridge(
+                    host=parser.searchGatewayResponse.ipaddress,
+                    uuid=uuid,
+                    bridge_type=parser.searchGatewayResponse.type,
+                )
             )
-        )
 
         # When we have passed a target, we only want to listen for that one
         if self._target:
@@ -65,26 +84,38 @@ def disconnect(self):
         """Disconnect the socket."""
         if self.transport:
             self.transport.close()
-        self._future.set_result(self._bridges)
+        if not self._future.done():
+            self._future.set_result(self._bridges)
 
     def get_bridges(self):
         """Return the discovered bridges."""
         return self._future
 
 
-async def discover_bridges(host: str = None, timeout: int = 1, loop=None) -> List[Bridge]:
-    """Discover a bridge by IP."""
+async def discover_bridges(host: str = None, timeout: int = 1, loop=None, broadcast_addresses: Iterable[Any] = None) -> List[Bridge]:
+    """Discover bridges on the network, or by IP.
+
+    The discovery request is sent to the limited broadcast address by default, which only reaches the
+    interface that the routing table picks. Pass the broadcast address of every network you want to
+    search in broadcast_addresses to reach bridges behind the other interfaces as well. In Home
+    Assistant, `network.async_get_ipv4_broadcast_addresses()` provides these.
+    """
 
     if loop is None:
         loop = asyncio.get_event_loop()
 
     transport, protocol = await loop.create_datagram_endpoint(
-        lambda: BridgeDiscoveryProtocol(host, timeout),
+        lambda: BridgeDiscoveryProtocol(host, timeout, broadcast_addresses),
         local_addr=("0.0.0.0", 0),
         allow_broadcast=not host,
     )
 
     try:
-        return await protocol.get_bridges()
+        bridges = await protocol.get_bridges()
     finally:
         transport.close()
+
+    if not bridges:
+        _LOGGER.info("No bridges responded to the discovery request")
+
+    return bridges

tests/test_discovery.py

@@ -0,0 +1,113 @@
+"""Tests for the bridge discovery."""
+
+from ipaddress import IPv4Address
+from unittest.mock import MagicMock
+
+import pytest
+
+from aiocomfoconnect.bridge import Bridge
+from aiocomfoconnect.discovery import BROADCAST_ADDRESS, BridgeDiscoveryProtocol
+from aiocomfoconnect.protobuf import zehnder_pb2
+
+BRIDGE_UUID = "0000000000221111111111111111ffff"
+
+
+def discovery_response(host: str = "192.168.1.213", uuid: str = BRIDGE_UUID, gateway_type: int = 0) -> bytes:
+    """Build a discovery response, like a bridge would send."""
+    operation = zehnder_pb2.DiscoveryOperation()  # pylint: disable=no-member
+    operation.searchGatewayResponse.ipaddress = host
+    operation.searchGatewayResponse.uuid = bytes.fromhex(uuid)
+    operation.searchGatewayResponse.version = 1
+    operation.searchGatewayResponse.type = gateway_type
+
+    return operation.SerializeToString()
+
+
+def targets_of(protocol: BridgeDiscoveryProtocol) -> list:
+    """Connect a mocked transport and return the addresses that were sent to."""
+    transport = MagicMock()
+    protocol.connection_made(transport)
+
+    return [call.args[1][0] for call in transport.sendto.call_args_list]
+
+
+class TestBridgeDiscoveryProtocol:
+    """Test the BridgeDiscoveryProtocol class."""
+
+    @pytest.mark.asyncio
+    async def test_broadcasts_to_limited_broadcast_by_default(self):
+        """Test that we keep broadcasting to the limited broadcast address when nothing is passed."""
+        protocol = BridgeDiscoveryProtocol()
+
+        assert targets_of(protocol) == [BROADCAST_ADDRESS]
+
+    @pytest.mark.asyncio
+    async def test_broadcasts_to_every_broadcast_address(self):
+        """Test that every broadcast address we know about is searched."""
+        protocol = BridgeDiscoveryProtocol(broadcast_addresses=["192.168.1.255", "10.0.0.255"])
+
+        assert targets_of(protocol) == ["192.168.1.255", "10.0.0.255"]
+
+    @pytest.mark.asyncio
+    async def test_accepts_ipaddress_objects(self):
+        """Test that we accept the IPv4Address objects that Home Assistant hands out."""
+        protocol = BridgeDiscoveryProtocol(broadcast_addresses=[IPv4Address("192.168.1.255")])
+
+        assert targets_of(protocol) == ["192.168.1.255"]
+
+    @pytest.mark.asyncio
+    async def test_host_takes_precedence(self):
+        """Test that we only send to the host when one is passed."""
+        protocol = BridgeDiscoveryProtocol(target="192.168.1.213", broadcast_addresses=["192.168.1.255"])
+
+        assert targets_of(protocol) == ["192.168.1.213"]
+
+    @pytest.mark.asyncio
+    async def test_bridge_is_reported_once(self):
+        """Test that a bridge that replies to multiple requests is only reported once."""
+        protocol = BridgeDiscoveryProtocol(broadcast_addresses=["192.168.1.255", "255.255.255.255"])
+        targets_of(protocol)
+
+        protocol.datagram_received(discovery_response(), ("192.168.1.213", Bridge.PORT))
+        protocol.datagram_received(discovery_response(), ("192.168.1.213", Bridge.PORT))
+        protocol.disconnect()
+
+        bridges = await protocol.get_bridges()
+        assert len(bridges) == 1
+        assert bridges[0].host == "192.168.1.213"
+        assert bridges[0].uuid == BRIDGE_UUID
+
+    @pytest.mark.asyncio
+    async def test_multiple_bridges_are_reported(self):
+        """Test that bridges on different networks are all reported."""
+        protocol = BridgeDiscoveryProtocol(broadcast_addresses=["192.168.1.255", "10.0.0.255"])
+        targets_of(protocol)
+
+        protocol.datagram_received(discovery_response(host="192.168.1.213"), ("192.168.1.213", Bridge.PORT))
+        protocol.datagram_received(discovery_response(host="10.0.0.213", uuid="0000000000222222222222222222ffff"), ("10.0.0.213", Bridge.PORT))
+        protocol.disconnect()
+
+        bridges = await protocol.get_bridges()
+        assert [bridge.host for bridge in bridges] == ["192.168.1.213", "10.0.0.213"]
+
+    @pytest.mark.asyncio
+    async def test_disconnect_is_idempotent(self):
+        """Test that a late response after the timeout doesn't blow up on the completed future."""
+        protocol = BridgeDiscoveryProtocol(target="192.168.1.213")
+        targets_of(protocol)
+
+        protocol.disconnect()  # The timeout expires.
+        protocol.datagram_received(discovery_response(), ("192.168.1.213", Bridge.PORT))  # Disconnects a second time.
+
+        assert await protocol.get_bridges() == []
+
+    @pytest.mark.asyncio
+    async def test_discovery_request_is_ignored(self):
+        """Test that we ignore the discovery requests of other clients."""
+        protocol = BridgeDiscoveryProtocol()
+        targets_of(protocol)
+
+        protocol.datagram_received(b"\x0a\x00", ("192.168.1.10", Bridge.PORT))
+        protocol.disconnect()
+
+        assert await protocol.get_bridges() == []