generate_test_keys.py (6652B)
1 #!/usr/bin/env python3 2 3 # Copyright The Mbed TLS Contributors 4 # SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later 5 6 """Module generating EC and RSA keys to be used in test_suite_pk instead of 7 generating the required key at run time. This helps speeding up testing.""" 8 9 from typing import Iterator, List, Tuple 10 import re 11 import argparse 12 from mbedtls_framework.asymmetric_key_data import ASYMMETRIC_KEY_DATA 13 from mbedtls_framework.build_tree import guess_project_root 14 15 BYTES_PER_LINE = 16 16 17 def c_byte_array_literal_content(array_name: str, key_data: bytes) -> Iterator[str]: 18 yield 'const unsigned char ' 19 yield array_name 20 yield '[] = {' 21 for index in range(0, len(key_data), BYTES_PER_LINE): 22 yield '\n ' 23 for b in key_data[index:index + BYTES_PER_LINE]: 24 yield ' {:#04x},'.format(b) 25 yield '\n};' 26 27 def convert_der_to_c(array_name: str, key_data: bytes) -> str: 28 return ''.join(c_byte_array_literal_content(array_name, key_data)) 29 30 def get_key_type(key: str) -> str: 31 if re.match('PSA_KEY_TYPE_RSA_.*', key): 32 return "rsa" 33 elif re.match('PSA_KEY_TYPE_ECC_.*', key): 34 return "ec" 35 else: 36 print("Unhandled key type {}".format(key)) 37 return "unknown" 38 39 def get_ec_key_family(key: str) -> str: 40 match = re.search(r'.*\((.*)\)', key) 41 if match is None: 42 raise Exception("Unable to get EC family from {}".format(key)) 43 return match.group(1) 44 45 # Legacy EC group ID do not support all the key types that PSA does, so the 46 # following dictionaries are used for: 47 # - getting prefix/suffix for legacy curve names 48 # - understand if the curve is supported in legacy symbols (MBEDTLS_ECP_DP_...) 49 EC_NAME_CONVERSION = { 50 'PSA_ECC_FAMILY_SECP_K1': { 51 192: ('secp', 'k1'), 52 224: ('secp', 'k1'), 53 256: ('secp', 'k1') 54 }, 55 'PSA_ECC_FAMILY_SECP_R1': { 56 192: ('secp', 'r1'), 57 224: ('secp', 'r1'), 58 256: ('secp', 'r1'), 59 384: ('secp', 'r1'), 60 521: ('secp', 'r1') 61 }, 62 'PSA_ECC_FAMILY_BRAINPOOL_P_R1': { 63 256: ('bp', 'r1'), 64 384: ('bp', 'r1'), 65 512: ('bp', 'r1') 66 }, 67 'PSA_ECC_FAMILY_MONTGOMERY': { 68 255: ('curve', '19'), 69 448: ('curve', '') 70 } 71 } 72 73 def get_ec_curve_name(priv_key: str, bits: int) -> str: 74 ec_family = get_ec_key_family(priv_key) 75 try: 76 prefix = EC_NAME_CONVERSION[ec_family][bits][0] 77 suffix = EC_NAME_CONVERSION[ec_family][bits][1] 78 except KeyError: 79 return "" 80 return prefix + str(bits) + suffix 81 82 def get_look_up_table_entry(key_type: str, group_id_or_keybits: str, 83 priv_array_name: str, pub_array_name: str) -> Iterator[str]: 84 if key_type == "ec": 85 yield " {{ {}, 0,\n".format(group_id_or_keybits) 86 else: 87 yield " {{ 0, {},\n".format(group_id_or_keybits) 88 yield " {0}, sizeof({0}),\n".format(priv_array_name) 89 yield " {0}, sizeof({0}) }},".format(pub_array_name) 90 91 92 def write_output_file(output_file_name: str, arrays: str, look_up_table: str): 93 with open(output_file_name, 'wt') as output: 94 output.write("""\ 95 /********************************************************************************* 96 * This file was automatically generated from framework/scripts/generate_test_keys.py. 97 * Please do not edit it manually. 98 *********************************************************************************/ 99 """) 100 output.write(arrays) 101 output.write(""" 102 struct predefined_key_element {{ 103 int group_id; // EC group ID; 0 for RSA keys 104 int keybits; // bits size of RSA key; 0 for EC keys 105 const unsigned char *priv_key; 106 size_t priv_key_len; 107 const unsigned char *pub_key; 108 size_t pub_key_len; 109 }}; 110 111 struct predefined_key_element predefined_keys[] = {{ 112 {} 113 }}; 114 115 /* End of generated file */ 116 """.format(look_up_table)) 117 118 def collect_keys() -> Tuple[str, str]: 119 """" 120 This function reads key data from ASYMMETRIC_KEY_DATA and, only for the 121 keys supported in legacy ECP/RSA modules, it returns 2 strings: 122 - the 1st contains C arrays declaration of these keys and 123 - the 2nd contains the final look-up table for all these arrays. 124 """ 125 arrays = [] 126 look_up_table = [] 127 128 # Get a list of private keys only in order to get a single item for every 129 # (key type, key bits) pair. We know that ASYMMETRIC_KEY_DATA 130 # contains also the public counterpart. 131 priv_keys = [key for key in ASYMMETRIC_KEY_DATA if '_KEY_PAIR' in key] 132 priv_keys = sorted(priv_keys) 133 134 for priv_key in priv_keys: 135 key_type = get_key_type(priv_key) 136 # Ignore keys which are not EC or RSA 137 if key_type == "unknown": 138 continue 139 140 pub_key = re.sub('_KEY_PAIR', '_PUBLIC_KEY', priv_key) 141 142 for bits in ASYMMETRIC_KEY_DATA[priv_key]: 143 if key_type == "ec": 144 curve = get_ec_curve_name(priv_key, bits) 145 # Ignore EC curves unsupported in legacy symbols 146 if curve == "": 147 continue 148 # Create output array name 149 if key_type == "rsa": 150 array_name_base = "_".join(["test", key_type, str(bits)]) 151 else: 152 array_name_base = "_".join(["test", key_type, curve]) 153 array_name_priv = array_name_base + "_priv" 154 array_name_pub = array_name_base + "_pub" 155 # Convert bytearray to C array 156 c_array_priv = convert_der_to_c(array_name_priv, ASYMMETRIC_KEY_DATA[priv_key][bits]) 157 c_array_pub = convert_der_to_c(array_name_pub, ASYMMETRIC_KEY_DATA[pub_key][bits]) 158 # Write the C array to the output file 159 arrays.append(''.join(["\n", c_array_priv, "\n", c_array_pub, "\n"])) 160 # Update the lookup table 161 if key_type == "ec": 162 group_id_or_keybits = "MBEDTLS_ECP_DP_" + curve.upper() 163 else: 164 group_id_or_keybits = str(bits) 165 look_up_table.append(''.join(get_look_up_table_entry(key_type, group_id_or_keybits, 166 array_name_priv, array_name_pub))) 167 168 return ''.join(arrays), '\n'.join(look_up_table) 169 170 def main() -> None: 171 default_output_path = guess_project_root() + "/tests/include/test/test_keys.h" 172 173 argparser = argparse.ArgumentParser() 174 argparser.add_argument("--output", help="Output file", default=default_output_path) 175 args = argparser.parse_args() 176 177 output_file = args.output 178 179 arrays, look_up_table = collect_keys() 180 181 write_output_file(output_file, arrays, look_up_table) 182 183 if __name__ == '__main__': 184 main()