Updated SmartPlug commands - set_power() and check_power()

1 files changed, 14 insertions, 2 deletions

diff --git a/broadlink/__init__.py b/broadlink/__init__.py
index 2f3966c4762f..45c3d06011e0 100755
--- a/broadlink/__init__.py
+++ b/broadlink/__init__.py
@@ -181,11 +181,23 @@ class device:
     return response[0]
 
   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]))
+
   def send_data(self, data):
     packet = bytearray([0x02, 0x00, 0x00, 0x00])
     packet += data