From 52d7cda76abb994865f5d86218d4f6bef0586168 Mon Sep 17 00:00:00 2001
From: happyleaves <happyleaves.tfr@gmail.com>
Date: Tue, 15 Dec 2015 18:44:22 -0500
Subject: [PATCH] improved multiple-device support

---
 README.md     |  6 ++--
 orvibo/s20.py | 87 ++++++++++++++++++++++++++++++++++++++-------------
 setup.py      |  2 +-
 3 files changed, 70 insertions(+), 25 deletions(-)

diff --git a/README.md b/README.md
index 6d2cb7b..26e0977 100644
--- a/README.md
+++ b/README.md
@@ -5,9 +5,10 @@ Control Orvibo devices with Python 3. Currently supports the S20 WiFi Smart Swit
 ## Usage
 
 ```python
-from orvibo.s20 import S20
+from orvibo.s20 import S20, discover
 
-s20 = S20("x.x.x.x") # Discover the IP on your own.
+hosts = discover() # Discover devices on your local network.
+s20 = S20("x.x.x.x") # Use a discovered host, or supply a known host.
 print(s20.on) # Current state (True = ON, False = OFF).
 s20.on = True # Turn it on.
 s20.on = False # Turn it off.
@@ -17,7 +18,6 @@ s20.on = False # Turn it off.
 
 Pull requests are welcome. Possible areas for improvement:
 
-* Discover configured devices (get IPs).
 * Additional Orvibo devices.
 * Expand S20 functions: Timers, configuration, etc
 
diff --git a/orvibo/s20.py b/orvibo/s20.py
index d2d6dbc..5899846 100644
--- a/orvibo/s20.py
+++ b/orvibo/s20.py
@@ -3,6 +3,7 @@
 import binascii
 import logging
 import socket
+import threading
 import time
 
 _LOGGER = logging.getLogger(__name__)
@@ -15,6 +16,9 @@
 TIMEOUT = 1.0
 DISCOVERY_TIMEOUT = 1.0
 
+# Timeout after which to renew device subscriptions
+SUBSCRIPTION_TIMEOUT = 60
+
 # Packet constants.
 MAGIC = b'\x68\x64'
 DISCOVERY = b'\x00\x06\x71\x61'
@@ -28,8 +32,50 @@
 ON = b'\x01'
 OFF = b'\x00'
 
-# Timeout after which to renew device subscriptions
-SUBSCRIPTION_TIMEOUT = 60
+# Socket
+_SOCKET = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+
+# Buffer
+_BUFFER = {}
+
+
+def _listen():
+    """ Listen on socket. """
+    while True:
+        data, addr = _SOCKET.recvfrom(1024)
+        _BUFFER[addr[0]] = data
+
+
+def _setup():
+    """ Set up module.
+
+    Open a UDP socket, and listen in a thread.
+    """
+    for opt in [socket.SO_BROADCAST, socket.SO_REUSEADDR, socket.SO_REUSEPORT]:
+        _SOCKET.setsockopt(socket.SOL_SOCKET, opt, 1)
+    _SOCKET.bind(('', PORT))
+    udp = threading.Thread(target=_listen, daemon=True)
+    udp.start()
+
+
+def discover(timeout=DISCOVERY_TIMEOUT):
+    """ Discover devices on the local network.
+
+    :param timeout: Optional timeout in seconds.
+    :returns: Set of discovered host addresses.
+    """
+    hosts = set()
+    payload = MAGIC + DISCOVERY
+    for _ in range(RETRIES):
+        _SOCKET.sendto(bytearray(payload), ('255.255.255.255', PORT))
+        start = time.time()
+        while time.time() < start + timeout:
+            for host, data in _BUFFER.copy().items():
+                if _is_discovery_response(data):
+                    if host not in hosts:
+                        _LOGGER.debug("Discovered device at %s", host)
+                    hosts.add(host)
+    return hosts
 
 
 def _is_discovery_response(data):
@@ -74,13 +120,7 @@ def __init__(self, host):
         :param host: IP or hostname of device.
         """
         self.host = host
-        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        for opt in [socket.SO_BROADCAST, socket.SO_REUSEADDR,
-                    socket.SO_REUSEPORT]:
-            self._socket.setsockopt(socket.SOL_SOCKET, opt, 1)
-        self._socket.bind(('', PORT))
         (self._mac, self._mac_reversed) = self._discover_mac()
-
         self._subscribe()
 
     @property
@@ -144,6 +184,10 @@ def _subscribe(self):
                 "No status could be found for {}".format(self.host))
 
     def _subscription_is_recent(self):
+        """ Check if subscription occurred recently.
+
+        :returns: Yes (True) or no (False)
+        """
         return self.last_subscribed > time.time() - SUBSCRIPTION_TIMEOUT
 
     def _control(self, state):
@@ -219,24 +263,22 @@ def _udp_transact(self, payload, handler, *args,
         :param broadcast: Send a broadcast instead.
         :param timeout: Timeout in seconds.
         """
+        if self.host in _BUFFER:
+            del _BUFFER[self.host]
         host = self.host
         if broadcast:
             host = '255.255.255.255'
         retval = None
-        self._socket.settimeout(timeout)
         for _ in range(RETRIES):
-            self._socket.sendto(bytearray(payload), (host, PORT))
-            while True:
-                try:
-                    data, addr = self._socket.recvfrom(1024)
-                    # From the right device?
-                    if addr[0] == self.host:
-                        retval = handler(data, *args)
-                    # Return as soon as a response is received
-                    if retval:
-                        return retval
-                except socket.timeout:
-                    break
+            _SOCKET.sendto(bytearray(payload), (host, PORT))
+            start = time.time()
+            while time.time() < start + timeout:
+                data = _BUFFER.get(self.host, None)
+                if data:
+                    retval = handler(data, *args)
+                # Return as soon as a response is received
+                if retval:
+                    return retval
 
     def _turn_on(self):
         """ Turn on the device. """
@@ -245,3 +287,6 @@ def _turn_on(self):
     def _turn_off(self):
         """ Turn off the device. """
         self._control(OFF)
+
+
+_setup()
diff --git a/setup.py b/setup.py
index be1e9cd..383430e 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name='orvibo',
-    version='1.0.1',
+    version='1.1.0',
     description='Control Orvibo products.',
     url='https://github.com/happyleavesaoc/python-orvibo/',
     license='MIT',