Merge branch 'bugfix/ds_gen_test_cases_fix' into 'master'

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

Closes IDFGH-6601

See merge request espressif/esp-idf!16915

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,14 +55,14 @@ 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 << 4096) for x in range(NUM_MESSAGES)]
+    messages = [random.randrange(0, 1 << max_key_size) for x in range(NUM_MESSAGES)]
 
     with open('digital_signature_test_cases.h', 'w') as f:
         f.write('/*\n')
@@ -73,7 +83,7 @@ with open('digital_signature_test_cases.h', 'w') as f:
 
         # 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')
+        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,\n' % number_as_bignum_words(m))
@@ -92,7 +102,7 @@ with open('digital_signature_test_cases.h', 'w') as f:
             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]
+            sizes = supported_key_size[target]
             key_size = sizes[case % len(sizes)]
 
             private_key = rsa.generate_private_key(
@@ -120,27 +130,29 @@ with open('digital_signature_test_cases.h', 'w') as f:
             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) + \
+            # 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
-        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) + \
+            # 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
 
-        assert len(p) == 12672 / 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()
@@ -159,3 +171,17 @@ with open('digital_signature_test_cases.h', 'w') as f:
             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)