From a6d515b5cb81c8dab4c983620a1ee1e9d6791f3e Mon Sep 17 00:00:00 2001
From: Anders Melchiorsen <amelchio@nogoto.net>
Date: Wed, 8 Aug 2018 16:52:03 +0200
Subject: [PATCH] Support Unicode strings in encrypt/decrypt (#125)

---
 pyHS100/protocol.py            | 28 ++++++++++++++--------------
 pyHS100/tests/test_protocol.py | 20 ++++++++++++++++++++
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/pyHS100/protocol.py b/pyHS100/protocol.py
index a950418f..9b669e14 100644
--- a/pyHS100/protocol.py
+++ b/pyHS100/protocol.py
@@ -86,14 +86,16 @@ class TPLinkSmartHomeProtocol:
         :return: ciphertext request
         """
         key = TPLinkSmartHomeProtocol.INITIALIZATION_VECTOR
-        buffer = bytearray(struct.pack(">I", len(request)))
 
-        for char in request:
-            cipher = key ^ ord(char)
-            key = cipher
-            buffer.append(cipher)
+        plainbytes = request.encode()
+        buffer = bytearray(struct.pack(">I", len(plainbytes)))
 
-        return buffer
+        for plainbyte in plainbytes:
+            cipherbyte = key ^ plainbyte
+            key = cipherbyte
+            buffer.append(cipherbyte)
+
+        return bytes(buffer)
 
     @staticmethod
     def decrypt(ciphertext: bytes) -> str:
@@ -106,13 +108,11 @@ class TPLinkSmartHomeProtocol:
         key = TPLinkSmartHomeProtocol.INITIALIZATION_VECTOR
         buffer = []
 
-        ciphertext_str = ciphertext.decode('latin-1')
+        for cipherbyte in ciphertext:
+            plainbyte = key ^ cipherbyte
+            key = cipherbyte
+            buffer.append(plainbyte)
 
-        for char in ciphertext_str:
-            plain = key ^ ord(char)
-            key = ord(char)
-            buffer.append(chr(plain))
+        plaintext = bytes(buffer)
 
-        plaintext = ''.join(buffer)
-
-        return plaintext
+        return plaintext.decode()
diff --git a/pyHS100/tests/test_protocol.py b/pyHS100/tests/test_protocol.py
index b7899730..773c9e43 100644
--- a/pyHS100/tests/test_protocol.py
+++ b/pyHS100/tests/test_protocol.py
@@ -10,3 +10,23 @@ class TestTPLinkSmartHomeProtocol(TestCase):
         # encrypt adds a 4 byte header
         encrypted = encrypted[4:]
         self.assertEqual(d, TPLinkSmartHomeProtocol.decrypt(encrypted))
+
+    def test_encrypt_unicode(self):
+        d = "{'snowman': '\u2603'}"
+
+        e = bytes([208, 247, 132, 234, 133, 242, 159, 254, 144, 183,
+                   141, 173, 138, 104, 240, 115, 84, 41])
+
+        encrypted = TPLinkSmartHomeProtocol.encrypt(d)
+        # encrypt adds a 4 byte header
+        encrypted = encrypted[4:]
+
+        self.assertEqual(e, encrypted)
+
+    def test_decrypt_unicode(self):
+        e = bytes([208, 247, 132, 234, 133, 242, 159, 254, 144, 183,
+                   141, 173, 138, 104, 240, 115, 84, 41])
+
+        d = "{'snowman': '\u2603'}"
+
+        self.assertEqual(d, TPLinkSmartHomeProtocol.decrypt(e))