Merge pull request mjg59#16 from PeWu/smartplug

Matthew Garrett · web-flow · commit 7ac243b83837 · 2016-11-19T16:47:28.000-08:00
Updated SmartPlug commands - set_power() and check_power()
diff --git a/README.md b/README.md
@@ -44,7 +44,12 @@ Obtain sensor data from an A1:
 data = devices[0].check_sensors()
 ```
 
-Set power state on an SP2/SP3 (0 for off, 1 for on):
+Set power state on a SmartPlug SP2/SP3:
+```
+devices[0].set_power(True)
+```
+
+Check power state on a SmartPlug:
+```
+state = devices[0].check_power()
 ```
-devices[0].set_power(1)
-```
diff --git a/broadlink/__init__.py b/broadlink/__init__.py
@@ -241,11 +241,22 @@ def __init__ (self, host, mac):
     device.__init__(self, host, mac)
 
   def set_power(self, state):
-    packet = bytearray(8)
+    """Sets the power state of the smart plug."""
+    packet = bytearray(16)
     packet[0] = 2
-    packet[4] = state
+    packet[4] = 1 if state else 0
     self.send_packet(0x6a, packet)
 
+  def check_power(self):
+    """Returns the power state of the smart plug."""
+    packet = bytearray(16)
+    packet[0] = 1
+    response = self.send_packet(0x6a, packet)
+    err = ord(response[0x22]) | (ord(response[0x23]) << 8)
+    if err == 0:
+      aes = AES.new(str(self.key), AES.MODE_CBC, str(self.iv))
+      payload = aes.decrypt(str(response[0x38:]))
+      return bool(ord(payload[0x4]))
 
 class a1(device):
   def __init__ (self, host, mac):
