diff --git a/espsecure.py b/espsecure.py index 79f2fa317..af79ac227 100755 --- a/espsecure.py +++ b/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): diff --git a/test/test_espsecure.py b/test/test_espsecure.py index d093f82e9..5e01ed204 100755 --- a/test/test_espsecure.py +++ b/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__)))