espsecure: Fix flash_encrypt_data / flash_decrypt_data on Python 3

Closes #369

espsecure.py

@@ -272,7 +272,12 @@ def _flash_encryption_tweak_key(key, offset, tweak_range):
 
     Return tweaked key
     """
-    key = [ord(k) for k in key]
+    if esptool.PYTHON2:
+        key = [ord(k) for k in key]
+    else:
+        key = list(key)
+    assert len(key) == 32
+
     offset_bits = [(offset & (1 << x)) != 0 for x in range(24)]
 
     for bit in tweak_range:
@@ -281,7 +286,10 @@ def _flash_encryption_tweak_key(key, offset, tweak_range):
             # to how it is looked up in the tweak pattern table
             key[bit // 8] ^= 1 << (7 - (bit % 8))
 
-    return b"".join(chr(k) for k in key)
+    if esptool.PYTHON2:
+        return b"".join(chr(k) for k in key)
+    else:
+        return bytes(key)
 
 
 def generate_flash_encryption_key(args):

test/test_espsecure.py

@@ -10,6 +10,7 @@
 import io
 import sys
 import tempfile
+import zlib
 from collections import namedtuple
 
 TEST_DIR = os.path.abspath(os.path.dirname(__file__))
@@ -81,6 +82,7 @@ def test_digest_bootloader(self):
         finally:
             os.unlink(output_file.name)
 
+
 class ECDSASigningTests(EspSecureTestCase):
 
     VerifyArgs = namedtuple('verify_signature_args', [
@@ -170,6 +172,55 @@ def test_extract_binary_public_key(self):
             os.unlink(pub_keyfile2.name)
 
 
+class ESP32FlashEncryptionTests(EspSecureTestCase):
+
+    def test_encrypt_decrypt(self):
+        EncryptArgs = namedtuple('encrypt_flash_data_args',
+                                 [ 'keyfile',
+                                   'output',
+                                   'address',
+                                   'flash_crypt_conf',
+                                   'plaintext_file'
+                                 ])
+
+        DecryptArgs = namedtuple('decrypt_flash_data_args',
+                                 [ 'keyfile',
+                                   'output',
+                                   'address',
+                                   'flash_crypt_conf',
+                                   'encrypted_file'
+                                 ])
+
+        original_plaintext = self._open('bootloader.bin')
+        keyfile = self._open('256bit_key.bin')
+        ciphertext = io.BytesIO()
+
+        args = EncryptArgs(keyfile,
+                           ciphertext,
+                           0x1000,
+                           0xFF,
+                           original_plaintext)
+        espsecure.encrypt_flash_data(args)
+
+        self.assertNotEqual(original_plaintext, ciphertext.getvalue())
+        # use compressed size as entropy estimate for effectiveness
+        compressed_cipher = zlib.compress(ciphertext.getvalue())
+        self.assertGreaterEqual(len(compressed_cipher), len(ciphertext.getvalue()))
+
+        ciphertext.seek(0)
+        keyfile.seek(0)
+        plaintext = io.BytesIO()
+        args = DecryptArgs(keyfile,
+                           plaintext,
+                           0x1000,
+                           0xFF,
+                           ciphertext)
+        espsecure.decrypt_flash_data(args)
+
+        original_plaintext.seek(0)
+        self.assertEqual(original_plaintext.read(), plaintext.getvalue())
+
+
 if __name__ == '__main__':
     print("Running espsecure tests...")
     print("Using espsecure %s at %s" % (esptool.__version__, os.path.abspath(espsecure.__file__)))