From 997feb74cc0fde004898f23da4bd93380b519b1b Mon Sep 17 00:00:00 2001
From: Brett Williams <blw1138@mac.com>
Date: Thu, 1 Jun 2023 18:09:26 -0500
Subject: [PATCH] Add zeroconf to client

---
 lib/client/client.py          | 33 +++++++++++------
 lib/server/zeroconf_server.py | 70 +++++++++++++++++++++++++++++++++++
 requirements.txt              |  3 +-
 3 files changed, 93 insertions(+), 13 deletions(-)
 create mode 100644 lib/server/zeroconf_server.py

diff --git a/lib/client/client.py b/lib/client/client.py
index 21031b9..71adccf 100644
--- a/lib/client/client.py
+++ b/lib/client/client.py
@@ -10,6 +10,7 @@ from tkinter import ttk, messagebox
 from PIL import Image, ImageTk
 from new_job_window import NewJobWindow
 from server_proxy import RenderServerProxy
+from lib.server.zeroconf_server import ZeroconfServer
 
 
 def request_data(server_ip, payload, server_port=8080, timeout=2):
@@ -44,15 +45,17 @@ class ZordonClient:
 
     def __init__(self):
 
-        servers = available_servers()
-
         # Create a Treeview widget
         self.root = tk.Tk()
         self.root.title("Zordon Render Client")
         self.local_host = socket.gethostname()
-        self.server_proxy = RenderServerProxy(hostname=servers[0])
+        self.server_proxy = RenderServerProxy(hostname=self.local_host)
         self.job_cache = []
 
+        # Setup zeroconf
+        self.zeroconf = ZeroconfServer("_zordon._tcp.local.", self.server_proxy.hostname, self.server_proxy.port)
+        self.zeroconf.start()
+
         # Setup photo preview
         photo_pad = tk.Frame(self.root, background="gray")
         photo_pad.pack(fill=tk.BOTH, pady=5, padx=5)
@@ -129,7 +132,7 @@ class ZordonClient:
             pass
 
         # update servers
-        self.populate_server_tree()
+        self.update_servers()
         try:
             selected_server = self.server_tree.get_children()[0]
             self.server_tree.selection_set(selected_server)
@@ -141,15 +144,12 @@ class ZordonClient:
         x.daemon = True
         x.start()
 
-    def populate_server_tree(self):
-        servers = available_servers()
-        self.server_tree.delete(*self.server_tree.get_children())
-        for hostname in servers:
-            self.server_tree.insert("", tk.END, iid=hostname, values=(hostname,))
-
     def server_picked(self, event):
-        new_hostname = self.server_tree.selection()[0]
-        self.server_proxy.hostname = new_hostname
+        try:
+            new_hostname = self.server_tree.selection()[0]
+            self.server_proxy.hostname = new_hostname
+        except IndexError:
+            pass
         self.job_cache.clear()
         self.update_jobs(clear_table=True)
 
@@ -218,8 +218,17 @@ class ZordonClient:
     def __background_update(self):
         while True:
             self.update_jobs()
+            self.update_servers()
             time.sleep(1)
 
+    def update_servers(self):
+        servers = self.zeroconf.found_clients()
+        if len(servers) < len(self.server_tree.get_children()):
+            self.server_tree.delete(*self.server_tree.get_children())
+        for hostname in servers:
+            if hostname not in self.server_tree.get_children():
+                self.server_tree.insert("", tk.END, iid=hostname, values=(hostname,))
+
     def update_jobs(self, clear_table=False):
 
         def update_row(tree, id, new_values, tags=None):
diff --git a/lib/server/zeroconf_server.py b/lib/server/zeroconf_server.py
new file mode 100644
index 0000000..2167c14
--- /dev/null
+++ b/lib/server/zeroconf_server.py
@@ -0,0 +1,70 @@
+import logging
+import socket
+import time
+
+from zeroconf import Zeroconf, ServiceInfo, ServiceBrowser, ServiceStateChange
+
+logger = logging.getLogger()
+
+
+class ZeroconfServer:
+    def __init__(self, service_type, server_name, server_port):
+        self.service_type = service_type
+        self.server_name = server_name
+        self.server_port = server_port
+        self.server_ip = None
+        self.zeroconf = Zeroconf()
+        self.service_info = None
+        self.client_cache = {}
+
+    def start(self):
+        self._register_service()
+        self._browse_services()
+
+    def stop(self):
+        self._unregister_service()
+        self.zeroconf.close()
+
+    def _register_service(self):
+        self.server_ip = socket.gethostbyname(socket.gethostname())
+
+        info = ServiceInfo(
+            self.service_type,
+            f"{self.server_name}.{self.service_type}",
+            addresses=[socket.inet_aton(self.server_ip)],
+            port=self.server_port,
+            properties={},
+        )
+
+        self.service_info = info
+        self.zeroconf.register_service(info)
+
+    def _unregister_service(self):
+        if self.service_info:
+            self.zeroconf.unregister_service(self.service_info)
+            self.service_info = None
+
+    def _browse_services(self):
+        browser = ServiceBrowser(self.zeroconf, self.service_type, [self._on_service_discovered])
+
+    def _on_service_discovered(self, zeroconf, service_type, name, state_change=None):
+        info = zeroconf.get_service_info(service_type, name)
+        if service_type == self.service_type:
+            if state_change == ServiceStateChange.Added or state_change == ServiceStateChange.Updated:
+                self.client_cache[name] = info
+            else:
+                self.client_cache.pop(name)
+
+    def found_clients(self):
+        return [x.split(f'.{self.service_type}')[0] for x in self.client_cache.keys()]
+
+
+# Example usage:
+if __name__ == "__main__":
+    server = ZeroconfServer("_zordon._tcp.local.", "Zordon", 5000)
+    server.start()
+
+    # Run your Flask application or perform other tasks
+    # ...
+    input("Server running - Press enter to end")
+    server.stop()
diff --git a/requirements.txt b/requirements.txt
index 868b63a..48d1a88 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,4 +10,5 @@ tkinterdnd2~=0.3.0
 future==0.18.3
 json2html~=1.3.0
 SQLAlchemy~=2.0.15
-Pillow~=9.3.0
\ No newline at end of file
+Pillow~=9.3.0
+zeroconf~=0.63.0
\ No newline at end of file