From 447f0aa9d4e4c1c67d76d763996236a797c252ae Mon Sep 17 00:00:00 2001
From: Harry Tung <honeytung@gmail.com>
Date: Wed, 24 Jul 2024 10:12:10 -0700
Subject: [PATCH] Added Workaround for WSDiscovery Bug (#49)

* Added exceptions to handle closing WSD properly

* Moved wsd to class variable

* Changed WSD to be singleton

* Added docstring

* Bumped version

* Automatically reformatting code with black and isort

* Added comments

---------

Co-authored-by: Auto-format Bot <autoformatbot@groundlight.ai>
---
 pyproject.toml                  |  2 +-
 src/framegrab/rtsp_discovery.py | 52 ++++++++++++++++++++++-----------
 2 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index f853505..8c29fe6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "framegrab"
-version = "0.6.0"
+version = "0.6.1"
 description = "Easily grab frames from cameras or streams"
 authors = ["Groundlight <info@groundlight.ai>"]
 license = "MIT"
diff --git a/src/framegrab/rtsp_discovery.py b/src/framegrab/rtsp_discovery.py
index 4491625..a74cd4d 100644
--- a/src/framegrab/rtsp_discovery.py
+++ b/src/framegrab/rtsp_discovery.py
@@ -67,6 +67,21 @@ class ONVIFDeviceInfo(BaseModel):
 class RTSPDiscovery:
     """Simple RTSP camera discovery with ONVIF capabilities"""
 
+    _wsd_instance = None
+
+    @classmethod
+    def _get_wsd(cls):
+        """
+        Get the WSDiscovery instance, creating it if it doesn't exist.
+
+        Returns:
+        WSDiscovery: The WSDiscovery instance.
+        """
+
+        if cls._wsd_instance is None:
+            cls._wsd_instance = WSDiscovery()
+        return cls._wsd_instance
+
     @staticmethod
     def discover_onvif_devices(
         auto_discover_mode: AutodiscoverMode = AutodiscoverMode.ip_only,
@@ -95,24 +110,27 @@ def discover_onvif_devices(
             logger.debug("ONVIF device discovery disabled")
             return device_ips
 
-        wsd = WSDiscovery()
-        wsd.start()
-        types = [QName("http://www.onvif.org/ver10/network/wsdl", "NetworkVideoTransmitter")]
-        ret = wsd.searchServices(types=types)
-        for service in ret:
-            xaddr = service.getXAddrs()[0]
-            parsed_url = urllib.parse.urlparse(xaddr)
-            ip = parsed_url.hostname
-            port = parsed_url.port or 80  # Use the default port 80 if not specified
-
-            logger.debug(f"Found ONVIF service at {xaddr}")
-            device_ip = ONVIFDeviceInfo(ip=ip, port=port, username="", password="", xaddr=xaddr, rtsp_urls=[])
-
-            if auto_discover_mode is not AutodiscoverMode.ip_only:
-                RTSPDiscovery._try_logins(device=device_ip, auto_discover_mode=auto_discover_mode)
+        try:
+            wsd = RTSPDiscovery._get_wsd()
+            wsd.start()
+            types = [QName("http://www.onvif.org/ver10/network/wsdl", "NetworkVideoTransmitter")]
+            ret = wsd.searchServices(types=types)
+            for service in ret:
+                xaddr = service.getXAddrs()[0]
+                parsed_url = urllib.parse.urlparse(xaddr)
+                ip = parsed_url.hostname
+                port = parsed_url.port or 80  # Use the default port 80 if not specified
+
+                logger.debug(f"Found ONVIF service at {xaddr}")
+                device_ip = ONVIFDeviceInfo(ip=ip, port=port, username="", password="", xaddr=xaddr, rtsp_urls=[])
+
+                if auto_discover_mode is not AutodiscoverMode.ip_only:
+                    RTSPDiscovery._try_logins(device=device_ip, auto_discover_mode=auto_discover_mode)
+
+                device_ips.append(device_ip)
+        finally:
+            wsd.stop()  # This is supposed to clean up the threads but it doesn't seem to work, sock is still open after running this
 
-            device_ips.append(device_ip)
-        wsd.stop()
         return device_ips
 
     @staticmethod