From a786694e932d48a6f194be04cfd47dff37611999 Mon Sep 17 00:00:00 2001 From: Jaxc Date: Fri, 9 Aug 2024 15:27:43 +0200 Subject: [PATCH] Add support for broadcast to all interfaces Also adds support for setting both own IP and broadcast IP --- PyStageLinQ/PyStageLinQ.py | 77 +++++++++++++------ README.md | 2 + tests/Main.py | 5 +- tests/unit/test_unit_PyStageLinQ.py | 114 +++++++++++++++++++++++----- 4 files changed, 155 insertions(+), 43 deletions(-) diff --git a/PyStageLinQ/PyStageLinQ.py b/PyStageLinQ/PyStageLinQ.py index cbebe73..d9dd953 100644 --- a/PyStageLinQ/PyStageLinQ.py +++ b/PyStageLinQ/PyStageLinQ.py @@ -28,9 +28,8 @@ class PyStageLinQ: :param name: This is the name which PyStageLinQ will announce itself with on the StageLinQ protocol. If not set it defaults to "Hello StageLinQ World". - :param ip: This is the ip of the interface you want to bind the sockets to, e.g. your local ethernet IP. - - :param announce_ip: The IP to send StageLinQ announcements to, 255.255.255.255 will send to every network. + :param ip: This is the ip of the interface you want to bind the sockets to, e.g. your local ethernet IP. If set to + None all interfaces will be used. """ @@ -45,8 +44,7 @@ def __init__( [str, StageLinQDiscovery, EngineServices.ServiceHandle], None ], name: str = "Hello StageLinQ World", - ip="169.254.13.37", - announce_ip="255.255.255.255", + ip=None, ): self.name = name self.OwnToken = StageLinQToken() @@ -62,7 +60,14 @@ def __init__( self.device_list = Device.DeviceList() - self.ip = ip + if ip is None: + interfaces = socket.getaddrinfo( + host=socket.gethostname(), port=None, family=socket.AF_INET + ) + self.ip = [ip[-1][0] for ip in interfaces] + + else: + self.ip = [ip] self.tasks = set() self.found_services = [] @@ -76,8 +81,6 @@ def __init__( self.new_device_found_callback = new_device_found_callback - self.announce_ip = announce_ip - def start_standalone(self): """ Function for starting PyStageLinq in standalone mode. In this mode this function will not return, but @@ -115,22 +118,25 @@ def _announce_self(self): self._send_discovery_frame(discovery_frame) def _send_discovery_frame(self, discovery_frame): - with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as discovery_socket: - discovery_socket.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1) - try: - discovery_socket.sendto( - discovery_frame, (self.announce_ip, self.StageLinQ_discovery_port) - ) - except PermissionError: - raise PermissionError( - f"Cannot write to IP {self.announce_ip}, " - f"this error could be due to that there is no network interface set up with this IP range" - ) + for ip in self.ip: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as discovery_socket: + discovery_socket.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1) + discovery_socket.bind((ip, 0)) + try: + discovery_socket.sendto( + discovery_frame, + ("255.255.255.255", self.StageLinQ_discovery_port), + ) + except PermissionError: + raise PermissionError( + f"Cannot send message on interface {ip}, " + f"this error could be due to that there is no network interface set up with this IP range" + ) def get_loop_condition(self) -> bool: return self._loopcondition - async def _discover_stagelinq_device(self, timeout=10): + async def _discover_stagelinq_device(self, host_ip, timeout=10): """ This function is used to find StageLinQ device announcements. """ @@ -141,7 +147,7 @@ async def _discover_stagelinq_device(self, timeout=10): discover_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) try: discover_socket.bind( - (self.ip, self.StageLinQ_discovery_port) + (host_ip, self.StageLinQ_discovery_port) ) # bind socket StageLinQ interface except: # Cannot bind to socket, check if IP is correct and link is up @@ -155,7 +161,7 @@ async def _discover_stagelinq_device(self, timeout=10): data_available = select.select([discover_socket], [], [], 0) if data_available[0]: data, addr = discover_socket.recvfrom(discover_buffer_size) - ip = addr[0] + device_ip = addr[0] discovery_frame = StageLinQDiscovery() if PyStageLinQError.STAGELINQOK != discovery_frame.decode_frame(data): @@ -178,12 +184,14 @@ async def _discover_stagelinq_device(self, timeout=10): ) if device_registered is True: continue - await self._register_new_device(discovery_frame, ip) + await self._register_new_device(discovery_frame, device_ip) if time.time() > loop_timeout: # No devices found within timeout return PyStageLinQError.DISCOVERYTIMEOUT + print(f"No discovery frames found on {host_ip} last {timeout} seconds.") + async def _register_new_device(self, discovery_frame, ip): stagelinq_device = StageLinQService(ip, discovery_frame, self.OwnToken, None) service_tasks = await stagelinq_device.get_tasks() @@ -232,6 +240,7 @@ async def _subscribe_to_statemap( async def _start_stagelinq(self, standalone=False): # Start the initial tasks of the library self.tasks.add(asyncio.create_task(self._periodic_announcement())) + self.tasks.add(asyncio.create_task(self._py_stagelinq_strapper())) if standalone: @@ -257,7 +266,27 @@ async def _periodic_announcement(self): await asyncio.sleep(0.5) async def _py_stagelinq_strapper(self): - return await self._discover_stagelinq_device(timeout=2) + strapper_tasks = set() + print(f"Looking for discovery frames on {len(self.ip)} IP local IP addresses:") + + for ip in self.ip: + print(f"{ip}") + strapper_tasks.add( + asyncio.create_task(self._discover_stagelinq_device(ip, timeout=2)) + ) + + while self.get_loop_condition(): + all_tasks_done = True + for task in strapper_tasks.copy(): + all_tasks_done = all_tasks_done and task.done() + if task.done(): + if task.exception() is not None: + raise task.exception() + + if all_tasks_done: + print("Timeout occurred on all interfaces.") + return + await asyncio.sleep(1) def stop(self): self._stop() diff --git a/README.md b/README.md index dbcd341..e9612fc 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,8 @@ def state_map_data_print(data): # Example main function, starting PyStageLinQ. if __name__ == "__main__": global PrimeGo + + # Run PyStageLinQ on all available network interfaces PrimeGo = PyStageLinQ.PyStageLinQ(new_device_found_callback, name="Jaxcie StagelinQ") PrimeGo.start() ``` diff --git a/tests/Main.py b/tests/Main.py index 615fffd..07e599c 100644 --- a/tests/Main.py +++ b/tests/Main.py @@ -45,8 +45,11 @@ def state_map_data_print(data): def main(): global PrimeGo + """PrimeGo = PyStageLinQ.PyStageLinQ( + new_device_found_callback, name="Jaxcie StageLinQ", ip="169.254.13.37" + )""" PrimeGo = PyStageLinQ.PyStageLinQ( - new_device_found_callback, name="Jaxcie StageLinQ", ip="255.255.255.255" + new_device_found_callback, name="Jaxcie StageLinQ" ) PrimeGo.start_standalone() diff --git a/tests/unit/test_unit_PyStageLinQ.py b/tests/unit/test_unit_PyStageLinQ.py index cd40215..0272b75 100644 --- a/tests/unit/test_unit_PyStageLinQ.py +++ b/tests/unit/test_unit_PyStageLinQ.py @@ -41,7 +41,7 @@ def test_init_values(dummy_pystagelinq, dummy_ip): assert ( type(dummy_pystagelinq.device_list) is PyStageLinQ.PyStageLinQ.Device.DeviceList ) - assert dummy_pystagelinq.ip == dummy_ip + assert dummy_pystagelinq.ip == [dummy_ip] assert dummy_pystagelinq.tasks == set() @@ -60,6 +60,17 @@ def test_init_values(dummy_pystagelinq, dummy_ip): assert dummy_pystagelinq.new_device_found_callback is None +def test_init_values_ip_none(monkeypatch): + dummy_ips = [[["1,2,3,4"]], [["5.6.7.8"]]] + socket = MagicMock() + socket.getaddrinfo.return_value = dummy_ips + monkeypatch.setattr(PyStageLinQ.PyStageLinQ, "socket", socket) + + dummy_pystagelinq = PyStageLinQ.PyStageLinQ.PyStageLinQ(None, ip=None) + + assert dummy_pystagelinq.ip == [dummy_ips[0][0][0], dummy_ips[1][0][0]] + + def test_start_standalone(dummy_pystagelinq, monkeypatch): run_mock = Mock() start_stagelinq_mock = Mock() @@ -129,6 +140,7 @@ def test_send_discovery_frame(dummy_pystagelinq, monkeypatch): discovery_socket = MagicMock() socket.socket.side_effect = discovery_socket + socket.getaddrinfo.side_effect = [[["255.255.255.255"]]] monkeypatch.setattr(PyStageLinQ.PyStageLinQ, "socket", socket) @@ -140,15 +152,18 @@ def test_send_discovery_frame(dummy_pystagelinq, monkeypatch): ) discovery_socket.return_value.__enter__.return_value.sendto.assert_called_once_with( dummy_discovery_frame, - (dummy_pystagelinq.announce_ip, dummy_pystagelinq.StageLinQ_discovery_port), + ("255.255.255.255", dummy_pystagelinq.StageLinQ_discovery_port), ) -def test_send_discovery_frame_permission_error(dummy_pystagelinq, monkeypatch): +def test_send_discovery_frame_permission_error( + dummy_pystagelinq, monkeypatch, dummy_ip +): dummy_discovery_frame = "AAAA" socket = MagicMock() discovery_socket = MagicMock() + socket.getaddrinfo.side_effect = [[["255.255.255.255"]]] socket.socket.side_effect = discovery_socket discovery_socket.return_value.__enter__.return_value.sendto.side_effect = ( @@ -167,20 +182,22 @@ def test_send_discovery_frame_permission_error(dummy_pystagelinq, monkeypatch): assert ( exception.value.args[0] - == f"Cannot write to IP {dummy_pystagelinq.announce_ip}, this error could be due to that there is no network " + == f"Cannot send message on interface {dummy_ip}, this error could be due to that there is no network " f"interface set up with this IP range" ) @pytest.mark.asyncio -async def test_discover_stagelinq_device_bind_error(dummy_pystagelinq, monkeypatch): +async def test_discover_stagelinq_device_bind_error( + dummy_pystagelinq, monkeypatch, dummy_ip +): socket = MagicMock() monkeypatch.setattr(PyStageLinQ.PyStageLinQ, "socket", socket) socket.socket.return_value.bind.side_effect = Exception() assert ( - await dummy_pystagelinq._discover_stagelinq_device() + await dummy_pystagelinq._discover_stagelinq_device(dummy_ip) == PyStageLinQError.CANNOTBINDSOCKET ) @@ -194,7 +211,9 @@ def test_get_loop_condition(dummy_pystagelinq): @pytest.mark.asyncio -async def test_discover_stagelinq_check_initialization(dummy_pystagelinq, monkeypatch): +async def test_discover_stagelinq_check_initialization( + dummy_pystagelinq, monkeypatch, dummy_ip +): socket = MagicMock() monkeypatch.setattr(PyStageLinQ.PyStageLinQ, "socket", socket) @@ -203,16 +222,16 @@ async def test_discover_stagelinq_check_initialization(dummy_pystagelinq, monkey dummy_pystagelinq, "get_loop_condition", get_loop_condition_mock ) - await dummy_pystagelinq._discover_stagelinq_device() + await dummy_pystagelinq._discover_stagelinq_device(dummy_ip) socket.socket.assert_called_once_with(socket.AF_INET, socket.SOCK_DGRAM) socket.socket.return_value.bind.assert_called_once_with( - (dummy_pystagelinq.ip, dummy_pystagelinq.StageLinQ_discovery_port) + (dummy_pystagelinq.ip[0], dummy_pystagelinq.StageLinQ_discovery_port) ) socket.socket.return_value.setblocking.assert_called_once_with(False) @pytest.mark.asyncio -async def test_discover_stagelinq_timeout(dummy_pystagelinq, monkeypatch): +async def test_discover_stagelinq_timeout(dummy_pystagelinq, monkeypatch, dummy_ip): socket = MagicMock() monkeypatch.setattr(PyStageLinQ.PyStageLinQ, "socket", socket) @@ -235,7 +254,7 @@ async def test_discover_stagelinq_timeout(dummy_pystagelinq, monkeypatch): select_mock.select.side_effect = [[False]] assert ( - await dummy_pystagelinq._discover_stagelinq_device() + await dummy_pystagelinq._discover_stagelinq_device(dummy_ip) == PyStageLinQError.DISCOVERYTIMEOUT ) @@ -270,7 +289,7 @@ async def test_discover_stagelinq_bad_frame(dummy_pystagelinq, monkeypatch, dumm socket.socket.return_value.recvfrom.side_effect = [[None, [dummy_ip, dummy_port]]] decode_frame_mock.side_effect = [PyStageLinQError.INVALIDFRAME] - assert await dummy_pystagelinq._discover_stagelinq_device() is None + assert await dummy_pystagelinq._discover_stagelinq_device(dummy_ip) is None sleep_mock.assert_called_once_with(0.1) select_mock.select.assert_called_once_with([socket.socket.return_value], [], [], 0) @@ -309,7 +328,7 @@ class discovery_dummy: socket.socket.return_value.recvfrom.side_effect = [[None, [dummy_ip, 0]]] time_mock.time.side_effect = [0, 5, 11] - assert await dummy_pystagelinq._discover_stagelinq_device() is None + assert await dummy_pystagelinq._discover_stagelinq_device(dummy_ip) is None assert time_mock.time.call_count == 2 @@ -347,7 +366,7 @@ class discovery_dummy: socket.socket.return_value.recvfrom.side_effect = [[None, [dummy_ip, 0]]] time_mock.time.side_effect = [0, 5, 11] - assert await dummy_pystagelinq._discover_stagelinq_device() is None + assert await dummy_pystagelinq._discover_stagelinq_device(dummy_ip) is None assert time_mock.time.call_count == 2 @@ -390,7 +409,7 @@ class discovery_dummy: dummy_pystagelinq.device_list = MagicMock() dummy_pystagelinq.device_list.find_registered_device.side_effect = [True] - assert await dummy_pystagelinq._discover_stagelinq_device() is None + assert await dummy_pystagelinq._discover_stagelinq_device(dummy_ip) is None assert time_mock.time.call_count == 2 dummy_pystagelinq.device_list.find_registered_device.assert_called_once_with( @@ -437,7 +456,7 @@ class discovery_dummy: dummy_pystagelinq.device_list = MagicMock() dummy_pystagelinq.device_list.find_registered_device.side_effect = [False] - assert await dummy_pystagelinq._discover_stagelinq_device() is None + assert await dummy_pystagelinq._discover_stagelinq_device(dummy_ip) is None assert time_mock.time.call_count == 3 register_device_mock.assert_called_once() @@ -762,6 +781,9 @@ async def test_wait_for_exit_task_exception(dummy_pystagelinq, monkeypatch): sleep_mock = AsyncMock() monkeypatch.setattr(PyStageLinQ.PyStageLinQ.asyncio, "sleep", sleep_mock) + socket = MagicMock() + monkeypatch.setattr(PyStageLinQ.PyStageLinQ, "socket", socket) + tasks_mock.copy.side_effect = [[task_mock]] get_loop_condition_mock.side_effect = [True, False] task_mock.done.side_effect = [True] @@ -813,7 +835,25 @@ async def test_periodic_announcement(dummy_pystagelinq, monkeypatch): @pytest.mark.asyncio -async def test_py_stagelinq_strapper(dummy_pystagelinq, monkeypatch): +async def test_py_stagelinq_strapper(dummy_pystagelinq, monkeypatch, dummy_ip): + discover_device_mock = AsyncMock() + monkeypatch.setattr( + dummy_pystagelinq, "_discover_stagelinq_device", discover_device_mock + ) + + await dummy_pystagelinq._py_stagelinq_strapper() + + discover_device_mock.assert_called_once_with(dummy_ip, timeout=2) + + +@pytest.mark.asyncio +async def test_py_stagelinq_strapper_loop_condition_false( + dummy_pystagelinq, monkeypatch, dummy_ip +): + get_loop_condition_mock = Mock(side_effect=[False]) + monkeypatch.setattr( + dummy_pystagelinq, "get_loop_condition", get_loop_condition_mock + ) discover_device_mock = AsyncMock() monkeypatch.setattr( dummy_pystagelinq, "_discover_stagelinq_device", discover_device_mock @@ -821,7 +861,45 @@ async def test_py_stagelinq_strapper(dummy_pystagelinq, monkeypatch): await dummy_pystagelinq._py_stagelinq_strapper() - discover_device_mock.assert_called_once_with(timeout=2) + discover_device_mock.assert_called_once_with(dummy_ip, timeout=2) + + +@pytest.mark.asyncio +async def test_py_stagelinq_strapper_task_exception( + dummy_pystagelinq, monkeypatch, dummy_ip +): + + asyncio_create_task_mock = MagicMock() + + asyncio_create_task_mock.return_value = asyncio_create_task_mock + + asyncio_create_task_mock.done.return_value = True + asyncio_create_task_mock.exception.return_value = RuntimeError + + monkeypatch.setattr( + PyStageLinQ.PyStageLinQ.asyncio, "create_task", asyncio_create_task_mock + ) + + get_loop_condition_mock = Mock() + monkeypatch.setattr( + dummy_pystagelinq, "get_loop_condition", get_loop_condition_mock + ) + + sleep_mock = AsyncMock() + monkeypatch.setattr(PyStageLinQ.PyStageLinQ.asyncio, "sleep", sleep_mock) + + discover_device_mock = MagicMock() + monkeypatch.setattr( + dummy_pystagelinq, "_discover_stagelinq_device", discover_device_mock + ) + + get_loop_condition_mock.side_effect = [True, False] + + with pytest.raises(RuntimeError) as exception: + await dummy_pystagelinq._py_stagelinq_strapper() + assert get_loop_condition_mock.call_count == 1 + + assert exception.type is RuntimeError def test_stop(dummy_pystagelinq, monkeypatch):