From 30cffe13d7035637c95c772e79ee1ad07431b51b Mon Sep 17 00:00:00 2001
From: David Rapan <david@rapan.cz>
Date: Mon, 10 Feb 2025 16:13:10 +0100
Subject: [PATCH] refactor: Discovery reply tool

---
 tools/discovery.py       |  8 ++++----
 tools/discovery_reply.py | 43 ++++++++++++++++++++++++----------------
 2 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/tools/discovery.py b/tools/discovery.py
index fa705df6..236b1184 100644
--- a/tools/discovery.py
+++ b/tools/discovery.py
@@ -1,9 +1,10 @@
 import socket
 import asyncio
 
+DISCOVERY_IP = "255.255.255.255"
 DISCOVERY_PORT = 48899
-DISCOVERY_TIMEOUT = 1
 DISCOVERY_MESSAGE = ["WIFIKIT-214028-READ".encode(), "HF-A11ASSISTHREAD".encode()]
+DISCOVERY_TIMEOUT = 1
 
 class DiscoveryProtocol:
     def __init__(self, addresses):
@@ -13,7 +14,7 @@ def __init__(self, addresses):
 
     def connection_made(self, transport):
         self.transport = transport
-        print(f"DiscoveryProtocol: Send to {self.addresses}")
+        print(f"DiscoveryProtocol: Send {DISCOVERY_MESSAGE[0]} to {self.addresses}")
         for address in self.addresses:
             self.transport.sendto(DISCOVERY_MESSAGE[0], (address, DISCOVERY_PORT))
 
@@ -31,11 +32,10 @@ def connection_lost(self, _):
 
 async def main():
     loop = asyncio.get_running_loop()
-    addresses = ["255.255.255.255"]
     wait = True
 
     try:
-        transport, protocol = await loop.create_datagram_endpoint(lambda: DiscoveryProtocol(addresses), family = socket.AF_INET, allow_broadcast = True)
+        transport, protocol = await loop.create_datagram_endpoint(lambda: DiscoveryProtocol([DISCOVERY_IP]), family = socket.AF_INET, allow_broadcast = True)
         r = None
         while r is None or wait:
             r = await asyncio.wait_for(protocol.responses.get(), DISCOVERY_TIMEOUT)
diff --git a/tools/discovery_reply.py b/tools/discovery_reply.py
index c004ffda..df7bd3ea 100644
--- a/tools/discovery_reply.py
+++ b/tools/discovery_reply.py
@@ -1,23 +1,32 @@
 import socket
+import asyncio
 import netifaces
 
-if __name__ == '__main__':
+DISCOVERY_IP = "0.0.0.0"
+DISCOVERY_PORT = 48899
+DISCOVERY_MESSAGE = ["WIFIKIT-214028-READ".encode(), "HF-A11ASSISTHREAD".encode()]
+
+ifaces = netifaces.ifaddresses(netifaces.gateways()['default'][2][1])
+iface_inet = ifaces[netifaces.AF_INET][0]
+iface_link = ifaces[netifaces.AF_LINK][0]
+
+class DiscoveryProtocol:
+    def connection_made(self, transport):
+        self.transport = transport
 
-    ifaces = netifaces.ifaddresses(netifaces.gateways()['default'][2][1])
-    iface_inet = ifaces[netifaces.AF_INET][0]
-    iface_link = ifaces[netifaces.AF_LINK][0]
+    def datagram_received(self, data, addr):
+        if data == DISCOVERY_MESSAGE[0]:
+            print(f"DiscoveryProtocol: Received {data} from {addr}")
+            self.transport.sendto(f"{iface_inet["addr"]},{iface_link["addr"].replace(':', '').upper()},1234567890".encode(), addr)
 
-    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
-        s.settimeout(1)
-        s.bind(('', 48899))
+async def main():
+    loop = asyncio.get_running_loop()
+    transport, _ = await loop.create_datagram_endpoint(DiscoveryProtocol, local_addr = (DISCOVERY_IP, DISCOVERY_PORT), family = socket.AF_INET, allow_broadcast = True)
 
-        while True:
-            try:
-                m = s.recvfrom(1024)
-                d = m[0].decode()
-                if d == "WIFIKIT-214028-READ":
-                    print(f"{m[0]} from {m[1]}")
-                    data = f"{iface_inet["addr"]},{iface_link["addr"].replace(':', '').upper()},1234567890"
-                    s.sendto(data.encode(), m[1])
-            except (TimeoutError, socket.timeout):
-                continue
\ No newline at end of file
+    try:
+        await asyncio.sleep(3600)
+    finally:
+        transport.close()
+
+if __name__ == '__main__':
+    asyncio.run(main())
\ No newline at end of file