1from __future__ import print_function
2
3import binascii
4import os
5import sys
6from collections import namedtuple
7from io import BytesIO
8
9import ttfw_idf
10
11try:
12    import espsecure
13except ImportError:
14    idf_path = os.getenv('IDF_PATH')
15    if not idf_path or not os.path.exists(idf_path):
16        raise
17    sys.path.insert(0, os.path.join(idf_path, 'components', 'esptool_py', 'esptool'))
18    import espsecure
19
20
21# To prepare a test runner for this example:
22# 1. Generate zero flash encryption key:
23#   dd if=/dev/zero of=key.bin bs=1 count=32
24# 2.Burn Efuses:
25#   espefuse.py --do-not-confirm -p $ESPPORT burn_efuse FLASH_CRYPT_CONFIG 0xf
26#   espefuse.py --do-not-confirm -p $ESPPORT burn_efuse FLASH_CRYPT_CNT 0x1
27#   espefuse.py --do-not-confirm -p $ESPPORT burn_key flash_encryption key.bin
28@ttfw_idf.idf_example_test(env_tag='Example_Flash_Encryption')
29def test_examples_security_flash_encryption(env, extra_data):
30    dut = env.get_dut('flash_encryption', 'examples/security/flash_encryption', dut_class=ttfw_idf.ESP32DUT)
31    # start test
32    dut.start_app()
33
34    # calculate the expected ciphertext
35    flash_addr = dut.app.partition_table['storage']['offset']
36    plain_hex_str = '00 01 02 03 04 05 06 07  08 09 0a 0b 0c 0d 0e 0f'
37    plain_data = binascii.unhexlify(plain_hex_str.replace(' ', ''))
38
39    # Emulate espsecure encrypt_flash_data command
40    EncryptFlashDataArgs = namedtuple('EncryptFlashDataArgs', ['output', 'plaintext_file', 'address', 'keyfile', 'flash_crypt_conf', 'aes_xts'])
41    args = EncryptFlashDataArgs(BytesIO(), BytesIO(plain_data), flash_addr, BytesIO(b'\x00' * 32), 0xF, None)
42    espsecure.encrypt_flash_data(args)
43
44    expected_ciphertext = args.output.getvalue()
45    hex_ciphertext = binascii.hexlify(expected_ciphertext).decode('ascii')
46    expected_str = (' '.join(hex_ciphertext[i:i + 2] for i in range(0, 16, 2)) + '  ' +
47                    ' '.join(hex_ciphertext[i:i + 2] for i in range(16, 32, 2)))
48
49    lines = [
50        'FLASH_CRYPT_CNT eFuse value is 1',
51        'Flash encryption feature is enabled in DEVELOPMENT mode',
52        'with esp_partition_write',
53        plain_hex_str,
54        'with esp_partition_read',
55        plain_hex_str,
56        'with spi_flash_read',
57        expected_str,
58        # The status of NVS encryption for the "nvs" partition
59        'NVS partition "nvs" is encrypted.'
60    ]
61    for line in lines:
62        dut.expect(line, timeout=2)
63
64
65if __name__ == '__main__':
66    test_examples_security_flash_encryption()
67