ds: update gen_digital_signature_tests.py to handle different max key sizes

Max key size is now decided by target parameter, and related parameters are
no longer hard coded.

Closes https://github.com/espressif/esp-idf/issues/8243

components/esp_hw_support/test/gen_digital_signature_tests.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
-# SPDX-FileCopyrightText: 2021 Espressif Systems (Shanghai) CO LTD
+# SPDX-FileCopyrightText: 2021-2022 Espressif Systems (Shanghai) CO LTD
 # SPDX-License-Identifier: Apache-2.0
 
+import argparse
 import datetime
 import hashlib
 import hmac
@@ -15,6 +16,15 @@ from cryptography.hazmat.primitives.asymmetric.rsa import _modinv as modinv  # t
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 from cryptography.utils import int_to_bytes
 
+supported_targets = {'esp32s2', 'esp32c3', 'esp32s3'}
+supported_key_size = {'esp32s2':[4096, 3072, 2048, 1024],
+                      'esp32c3':[3072, 2048, 1024],
+                      'esp32s3':[4096, 3072, 2048, 1024]}
+
+NUM_HMAC_KEYS = 3
+NUM_MESSAGES = 10
+NUM_CASES = 6
+
 
 def number_as_bignum_words(number):  # type: (int) -> str
     """
@@ -45,117 +55,133 @@ def bytes_as_char_array(b):  # type: (bytes) -> str
     return '{ ' + ', '.join('0x%02x' % x for x in b) + ' }'
 
 
-NUM_HMAC_KEYS = 3
-NUM_MESSAGES = 10
-NUM_CASES = 6
+def generate_tests_cases(target):  # type: (str) -> None
+
+    max_key_size = max(supported_key_size[target])
+    print('Generating tests cases for {} (max key size = {})'.format(target, max_key_size))
+
+    hmac_keys = [os.urandom(32) for x in range(NUM_HMAC_KEYS)]
 
+    messages = [random.randrange(0, 1 << max_key_size) for x in range(NUM_MESSAGES)]
 
-hmac_keys = [os.urandom(32) for x in range(NUM_HMAC_KEYS)]
-
-messages = [random.randrange(0, 1 << 4096) for x in range(NUM_MESSAGES)]
-
-with open('digital_signature_test_cases.h', 'w') as f:
-    f.write('/*\n')
-    year = datetime.datetime.now().year
-    f.write(' * SPDX-FileCopyrightText: {year} Espressif Systems (Shanghai) CO LTD\n'.format(year=year))
-    f.write(' *\n')
-    f.write(' * SPDX-License-Identifier: Apache-2.0\n')
-    f.write(' *\n')
-    f.write(' * File generated by gen_digital_signature_tests.py\n')
-    f.write(' */\n')
-
-    # Write out HMAC keys
-    f.write('#define NUM_HMAC_KEYS %d\n\n' % NUM_HMAC_KEYS)
-    f.write('static const uint8_t test_hmac_keys[NUM_HMAC_KEYS][32] = {\n')
-    for h in hmac_keys:
-        f.write('     %s,\n' % bytes_as_char_array(h))
-    f.write('};\n\n')
-
-    # Write out messages
-    f.write('#define NUM_MESSAGES %d\n\n' % NUM_MESSAGES)
-    f.write('static const uint32_t test_messages[NUM_MESSAGES][4096/32] = {\n')
-    for m in messages:
-        f.write('        // Message %d\n' % messages.index(m))
-        f.write('        %s,\n' % number_as_bignum_words(m))
-    f.write('    };\n')
-    f.write('\n\n\n')
-
-    f.write('#define NUM_CASES %d\n\n' % NUM_CASES)
-    f.write('static const encrypt_testcase_t test_cases[NUM_CASES] = {\n')
-
-    for case in range(NUM_CASES):
-        f.write('    { /* Case %d */\n' % case)
-
-        iv = os.urandom(16)
-        f.write('        .iv = %s,\n' % (bytes_as_char_array(iv)))
-
-        hmac_key_idx = random.randrange(0, NUM_HMAC_KEYS)
-        aes_key = hmac.HMAC(hmac_keys[hmac_key_idx], b'\xFF' * 32, hashlib.sha256).digest()
-
-        sizes = [4096, 3072, 2048, 1024, 512]
-        key_size = sizes[case % len(sizes)]
-
-        private_key = rsa.generate_private_key(
-            public_exponent=65537,
-            key_size=key_size,
-            backend=default_backend())
-
-        priv_numbers = private_key.private_numbers()
-        pub_numbers = private_key.public_key().public_numbers()
-        Y = priv_numbers.d
-        M = pub_numbers.n
-
-        rr = 1 << (key_size * 2)
-        rinv = rr % pub_numbers.n
-        mprime = - modinv(M, 1 << 32)
-        mprime &= 0xFFFFFFFF
-        length = key_size // 32 - 1
-
-        f.write('        .p_data = {\n')
-        f.write('            .Y = %s,\n' % number_as_bignum_words(Y))
-        f.write('            .M = %s,\n' % number_as_bignum_words(M))
-        f.write('            .Rb = %s,\n' % number_as_bignum_words(rinv))
-        f.write('            .M_prime = 0x%08x,\n' % mprime)
-        f.write('            .length = %d, // %d bit\n' % (length, key_size))
-        f.write('        },\n')
-
-        # calculate MD from preceding values and IV
-
-        # Y4096 || M4096 || Rb4096 || M_prime32 || LENGTH32 || IV128
-        md_in = number_as_bytes(Y, 4096) + \
-            number_as_bytes(M, 4096) + \
-            number_as_bytes(rinv, 4096) + \
-            struct.pack('<II', mprime, length) + \
-            iv
-        assert len(md_in) == 12480 / 8
-        md = hashlib.sha256(md_in).digest()
-
-        # generate expected C value from P bitstring
-        #
-        # Y4096 || M4096 || Rb4096 || M_prime32 || LENGTH32 || MD256 || 0x08*8
-        p = number_as_bytes(Y, 4096) + \
-            number_as_bytes(M, 4096) + \
-            number_as_bytes(rinv, 4096) + \
-            md + \
-            struct.pack('<II', mprime, length) + \
-            b'\x08' * 8
-
-        assert len(p) == 12672 / 8
-
-        cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
-        encryptor = cipher.encryptor()
-        c = encryptor.update(p) + encryptor.finalize()
-
-        f.write('        .expected_c = %s,\n' % bytes_as_char_array(c))
-        f.write('        .hmac_key_idx = %d,\n' % (hmac_key_idx))
-
-        f.write('        // results of message array encrypted with these keys\n')
-        f.write('        .expected_results = {\n')
-        mask = (1 << key_size) - 1  # truncate messages if needed
+    with open('digital_signature_test_cases.h', 'w') as f:
+        f.write('/*\n')
+        year = datetime.datetime.now().year
+        f.write(' * SPDX-FileCopyrightText: {year} Espressif Systems (Shanghai) CO LTD\n'.format(year=year))
+        f.write(' *\n')
+        f.write(' * SPDX-License-Identifier: Apache-2.0\n')
+        f.write(' *\n')
+        f.write(' * File generated by gen_digital_signature_tests.py\n')
+        f.write(' */\n')
+
+        # Write out HMAC keys
+        f.write('#define NUM_HMAC_KEYS %d\n\n' % NUM_HMAC_KEYS)
+        f.write('static const uint8_t test_hmac_keys[NUM_HMAC_KEYS][32] = {\n')
+        for h in hmac_keys:
+            f.write('     %s,\n' % bytes_as_char_array(h))
+        f.write('};\n\n')
+
+        # Write out messages
+        f.write('#define NUM_MESSAGES %d\n\n' % NUM_MESSAGES)
+        f.write('static const uint32_t test_messages[NUM_MESSAGES][%d/32] = {\n' % max_key_size)
         for m in messages:
             f.write('        // Message %d\n' % messages.index(m))
-            f.write('      %s,' % (number_as_bignum_words(pow(m & mask, Y, M))))
-        f.write('     },\n')
-        f.write('     },\n')
-
-    f.write('};\n')
+            f.write('        %s,\n' % number_as_bignum_words(m))
+        f.write('    };\n')
+        f.write('\n\n\n')
+
+        f.write('#define NUM_CASES %d\n\n' % NUM_CASES)
+        f.write('static const encrypt_testcase_t test_cases[NUM_CASES] = {\n')
+
+        for case in range(NUM_CASES):
+            f.write('    { /* Case %d */\n' % case)
+
+            iv = os.urandom(16)
+            f.write('        .iv = %s,\n' % (bytes_as_char_array(iv)))
+
+            hmac_key_idx = random.randrange(0, NUM_HMAC_KEYS)
+            aes_key = hmac.HMAC(hmac_keys[hmac_key_idx], b'\xFF' * 32, hashlib.sha256).digest()
+
+            sizes = supported_key_size[target]
+            key_size = sizes[case % len(sizes)]
+
+            private_key = rsa.generate_private_key(
+                public_exponent=65537,
+                key_size=key_size,
+                backend=default_backend())
+
+            priv_numbers = private_key.private_numbers()
+            pub_numbers = private_key.public_key().public_numbers()
+            Y = priv_numbers.d
+            M = pub_numbers.n
+
+            rr = 1 << (key_size * 2)
+            rinv = rr % pub_numbers.n
+            mprime = - modinv(M, 1 << 32)
+            mprime &= 0xFFFFFFFF
+            length = key_size // 32 - 1
+
+            f.write('        .p_data = {\n')
+            f.write('            .Y = %s,\n' % number_as_bignum_words(Y))
+            f.write('            .M = %s,\n' % number_as_bignum_words(M))
+            f.write('            .Rb = %s,\n' % number_as_bignum_words(rinv))
+            f.write('            .M_prime = 0x%08x,\n' % mprime)
+            f.write('            .length = %d, // %d bit\n' % (length, key_size))
+            f.write('        },\n')
+
+            # calculate MD from preceding values and IV
+            # Y_max_key_size || M_max_key_size || Rb_max_key_size || M_prime32 || LENGTH32 || IV128
+            md_in = number_as_bytes(Y, max_key_size) + \
+                number_as_bytes(M, max_key_size) + \
+                number_as_bytes(rinv, max_key_size) + \
+                struct.pack('<II', mprime, length) + \
+                iv
+
+            md = hashlib.sha256(md_in).digest()
+
+            # generate expected C value from P bitstring
+            #
+            # Y_max_key_size || M_max_key_size || Rb_max_key_size || M_prime32 || LENGTH32 ||  0x08*8
+            # E.g. for C3: Y3072 || M3072 || Rb3072 || M_prime32 || LENGTH32 || MD256 || 0x08*8
+            p = number_as_bytes(Y, max_key_size) + \
+                number_as_bytes(M, max_key_size) + \
+                number_as_bytes(rinv, max_key_size) + \
+                md + \
+                struct.pack('<II', mprime, length) + \
+                b'\x08' * 8
+
+            # expected_len = max_len_Y + max_len_M + max_len_rinv + md (32 bytes) + (mprime + length packed (8bytes)) + padding (8 bytes)
+            expected_len = (max_key_size / 8) * 3 + 32 + 8 + 8
+            assert len(p) == expected_len
+
+            cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
+            encryptor = cipher.encryptor()
+            c = encryptor.update(p) + encryptor.finalize()
+
+            f.write('        .expected_c = %s,\n' % bytes_as_char_array(c))
+            f.write('        .hmac_key_idx = %d,\n' % (hmac_key_idx))
+
+            f.write('        // results of message array encrypted with these keys\n')
+            f.write('        .expected_results = {\n')
+            mask = (1 << key_size) - 1  # truncate messages if needed
+            for m in messages:
+                f.write('        // Message %d\n' % messages.index(m))
+                f.write('      %s,' % (number_as_bignum_words(pow(m & mask, Y, M))))
+            f.write('     },\n')
+            f.write('     },\n')
+
+        f.write('};\n')
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='''Generates Digital Signature Test Cases''')
+
+    parser.add_argument(
+        '--target',
+        required=True,
+        choices=supported_targets,
+        help='Target to generate test cases for, different targets support different max key length')
+
+    args = parser.parse_args()
+
+    generate_tests_cases(args.target)