Merge pull request #16 from PeWu/smartplug

Updated SmartPlug commands - set_power() and check_power()

1 files changed, 13 insertions, 2 deletions

diff --git a/broadlink/__init__.py b/broadlink/__init__.py
index 5d5dd4c3fbf2..942b017ffcf6 100755
--- a/broadlink/__init__.py
+++ b/broadlink/__init__.py
@@ -241,11 +241,22 @@ class sp2(device):
     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):