Updated SmartPlug commands - set_power() and check_power()

2 files changed, 22 insertions, 5 deletions

diff --git a/README.md b/README.md
index 3de1fa92e7ef..47a4691f51bd 100644
--- 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)
-```
\ No newline at end of file
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