1 files changed, 90 insertions, 60 deletions
diff --git a/broadlink/__init__.py b/broadlink/__init__.py
index 63ad76b38677..edd5f24bfede 100755
--- a/broadlink/__init__.py
+++ b/broadlink/__init__.py
@@ -1,74 +1,95 @@
 #!/usr/bin/python
 
 from datetime import datetime
-from socket import *
 from Crypto.Cipher import AES
 import time
 import random
+import socket
 
-class rm2:
-  def __init__(self):
-    self.count = random.randrange(0xffff) 
-    self.key = bytearray([0x09, 0x76, 0x28, 0x34, 0x3f, 0xe9, 0x9e, 0x23, 0x76, 0x5c, 0x15, 0x13, 0xac, 0xcf, 0x8b, 0x02])
-    self.iv = bytearray([0x56, 0x2e, 0x17, 0x99, 0x6d, 0x09, 0x3d, 0x28, 0xdd, 0xb3, 0xba, 0x69, 0x5a, 0x2e, 0x6f, 0x58])
-    s = socket(AF_INET, SOCK_DGRAM)
-    s.connect(('8.8.8.8', 0))  # connecting to a UDP address doesn't send packets
-    local_ip_address = s.getsockname()[0]
+def discover(timeout=None):
+  s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+  s.connect(('8.8.8.8', 0))  # connecting to a UDP address doesn't send packets
+  local_ip_address = s.getsockname()[0]
+  address = local_ip_address.split('.')
+  cs = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+  cs.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+  cs.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
+  cs.bind(('',0))
+  port = cs.getsockname()[1]
+  starttime = time.time()
 
-    self.address = local_ip_address.split('.')
-    self.id = bytearray([0, 0, 0, 0])
+  devices = []
 
-  def discover(self):
-    self.cs = socket(AF_INET, SOCK_DGRAM)
-    self.cs.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
-    self.cs.setsockopt(SOL_SOCKET, SO_BROADCAST, 1)
-    self.cs.bind(('',0)) 
-    self.port = self.cs.getsockname()[1]
-
-    timezone = time.timezone/-3600
-    packet = bytearray(0x30)
-
-    year = datetime.now().year
-
-    if timezone < 0:
-      packet[0x08] = 0xff + timezone - 1
-      packet[0x09] = 0xff
-      packet[0x0a] = 0xff
-      packet[0x0b] = 0xff
-    else:
-      packet[0x08] = timezone
-      packet[0x09] = 0
-      packet[0x0a] = 0
-      packet[0x0b] = 0
-    packet[0x0c] = year & 0xff
-    packet[0x0d] = year >> 8
-    packet[0x0e] = datetime.now().minute
-    packet[0x0f] = datetime.now().hour
-    subyear = str(year)[2:]
-    packet[0x10] = int(subyear)
-    packet[0x11] = datetime.now().isoweekday()
-    packet[0x12] = datetime.now().day
-    packet[0x13] = datetime.now().month
-    packet[0x18] = int(self.address[0])
-    packet[0x19] = int(self.address[1])
-    packet[0x1a] = int(self.address[2])
-    packet[0x1b] = int(self.address[3])
-    packet[0x1c] = self.port & 0xff
-    packet[0x1d] = self.port >> 8
-    packet[0x26] = 6
-    checksum = 0xbeaf
+  timezone = time.timezone/-3600
+  packet = bytearray(0x30)
 
-    for i in range(len(packet)):
-        checksum += packet[i]
-    checksum = checksum & 0xffff
-    packet[0x20] = checksum & 0xff
-    packet[0x21] = checksum >> 8
+  year = datetime.now().year
 
-    self.cs.sendto(packet, ('255.255.255.255', 80))
-    response = self.cs.recvfrom(1024)
+  if timezone < 0:
+    packet[0x08] = 0xff + timezone - 1
+    packet[0x09] = 0xff
+    packet[0x0a] = 0xff
+    packet[0x0b] = 0xff
+  else:
+    packet[0x08] = timezone
+    packet[0x09] = 0
+    packet[0x0a] = 0
+    packet[0x0b] = 0
+  packet[0x0c] = year & 0xff
+  packet[0x0d] = year >> 8
+  packet[0x0e] = datetime.now().minute
+  packet[0x0f] = datetime.now().hour
+  subyear = str(year)[2:]
+  packet[0x10] = int(subyear)
+  packet[0x11] = datetime.now().isoweekday()
+  packet[0x12] = datetime.now().day
+  packet[0x13] = datetime.now().month
+  packet[0x18] = int(address[0])
+  packet[0x19] = int(address[1])
+  packet[0x1a] = int(address[2])
+  packet[0x1b] = int(address[3])
+  packet[0x1c] = port & 0xff
+  packet[0x1d] = port >> 8
+  packet[0x26] = 6
+  checksum = 0xbeaf
+
+  for i in range(len(packet)):
+      checksum += packet[i]
+  checksum = checksum & 0xffff
+  packet[0x20] = checksum & 0xff
+  packet[0x21] = checksum >> 8
+
+  cs.sendto(packet, ('255.255.255.255', 80))
+  if timeout is None:
+    response = cs.recvfrom(1024)
     responsepacket = bytearray(response[0])
-    self.host = response[1]
-    self.mac = responsepacket[0x3a:0x40]
+    host = response[1]
+    mac = responsepacket[0x3a:0x40]
+    return device(host=host, mac=mac)
+  else:
+    while (time.time() - starttime) < timeout:
+      cs.settimeout(timeout - (time.time() - starttime))
+      try:
+        response = cs.recvfrom(1024)
+      except socket.timeout:
+        return devices
+      responsepacket = bytearray(response[0])
+      host = response[1]
+      mac = responsepacket[0x3a:0x40]
+      devices.append(device(host=host, mac=mac))
+
+class device:
+  def __init__(self, host, mac):
+    self.host = host
+    self.mac = mac
+    self.count = random.randrange(0xffff)
+    self.key = bytearray([0x09, 0x76, 0x28, 0x34, 0x3f, 0xe9, 0x9e, 0x23, 0x76, 0x5c, 0x15, 0x13, 0xac, 0xcf, 0x8b, 0x02])
+    self.iv = bytearray([0x56, 0x2e, 0x17, 0x99, 0x6d, 0x09, 0x3d, 0x28, 0xdd, 0xb3, 0xba, 0x69, 0x5a, 0x2e, 0x6f, 0x58])
+    self.id = bytearray([0, 0, 0, 0])
+    self.cs = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    self.cs.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    self.cs.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
+    self.cs.bind(('',0))
 
   def auth(self):
     payload = bytearray(0x50)
@@ -102,7 +123,7 @@ class rm2:
     enc_payload = response[0x38:]
 
     aes = AES.new(str(self.key), AES.MODE_CBC, str(self.iv))
-    payload = aes.decrypt(str(response[0x38:]))
+    payload = aes.decrypt(str(enc_payload))
 
     self.id = payload[0x00:0x04]
     self.key = payload[0x04:0x14]
@@ -232,3 +253,12 @@ class rm2:
       aes = AES.new(str(self.key), AES.MODE_CBC, str(self.iv))
       payload = aes.decrypt(str(response[0x38:]))
       return payload[0x04:]
+
+class rm2(device):
+  def __init__ (self):
+    device.__init__(self, None, None)
+
+  def discover(self):
+    dev = discover()
+    self.host = dev.host
+    self.mac = dev.mac