diff --git a/aiocomfoconnect/__main__.py b/aiocomfoconnect/__main__.py index b6c4ee8..dc661e6 100644 --- a/aiocomfoconnect/__main__.py +++ b/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) diff --git a/aiocomfoconnect/discovery.py b/aiocomfoconnect/discovery.py index cf6453e..a923fc1 100644 --- a/aiocomfoconnect/discovery.py +++ b/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 = "" + 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", ("", 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 diff --git a/tests/test_discovery.py b/tests/test_discovery.py new file mode 100644 index 0000000..1401981 --- /dev/null +++ b/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() == []