diff --git a/tests/scripts/generate_test_keys.py b/tests/scripts/generate_test_keys.py index 65e427983a..335e84890b 100755 --- a/tests/scripts/generate_test_keys.py +++ b/tests/scripts/generate_test_keys.py @@ -6,7 +6,7 @@ """Module generating EC and RSA keys to be used in test_suite_pk instead of generating the required key at run time. This helps speeding up testing.""" -from typing import Iterator +from typing import Iterator, List import re import argparse import scripts_path # pylint: disable=unused-import @@ -89,8 +89,33 @@ def get_look_up_table_entry(key_type: str, group_id_or_keybits: str, yield " {0}, sizeof({0}),\n".format(priv_array_name) yield " {0}, sizeof({0}) }},".format(pub_array_name) + +def write_output_file(output_file_name: str, arrays: List[str], look_up_table: List[str]): + with open(output_file_name, 'wt') as output: + output.write(""" +/********************************************************************************* + * This file was automatically generated from tests/scripts/generate_test_keys.py. + * Please do not edit it manually. + *********************************************************************************/ +""") + output.write(''.join(arrays)) + output.write(""" +struct predefined_key_element {{ + int group_id; // EC group ID; 0 for RSA keys + int keybits; // bits size of RSA key; 0 for EC keys + const unsigned char *priv_key; + size_t priv_key_len; + const unsigned char *pub_key; + size_t pub_key_len; +}}; + +struct predefined_key_element predefined_keys[] = {{ +{} +}}; +""".format("\n".join(look_up_table))) + + def main() -> None: - #pylint: disable=too-many-locals default_output_path = guess_project_root() + "/tests/src/test_keys.h" argparser = argparse.ArgumentParser() @@ -98,14 +123,8 @@ def main() -> None: args = argparser.parse_args() output_file = args.output - output_file = open(output_file, 'wt') - output_file.write( - "/*********************************************************************************\n" + - " * This file was automatically generated from tests/scripts/generate_test_keys.py.\n" + - " * Please do not edit it manually.\n" + - " *********************************************************************************/\n" - ) + arrays = [] look_up_table = [] # Get a list of private keys only in order to get a single item for every @@ -139,7 +158,7 @@ def main() -> None: c_array_priv = convert_der_to_c(array_name_priv, ASYMMETRIC_KEY_DATA[priv_key][bits]) c_array_pub = convert_der_to_c(array_name_pub, ASYMMETRIC_KEY_DATA[pub_key][bits]) # Write the C array to the output file - output_file.write(''.join(["\n", c_array_priv, "\n", c_array_pub, "\n"])) + arrays.append(''.join(["\n", c_array_priv, "\n", c_array_pub, "\n"])) # Update the lookup table if key_type == "ec": group_id_or_keybits = "MBEDTLS_ECP_DP_" + curve.upper() @@ -147,22 +166,8 @@ def main() -> None: group_id_or_keybits = str(bits) look_up_table.append(''.join(get_look_up_table_entry(key_type, group_id_or_keybits, array_name_priv, array_name_pub))) - # Write the lookup table: the struct containing pointers to all the arrays we created above. - output_file.write(""" -struct predefined_key_element { - int group_id; // EC group ID; 0 for RSA keys - int keybits; // bits size of RSA key; 0 for EC keys - const unsigned char *priv_key; - size_t priv_key_len; - const unsigned char *pub_key; - size_t pub_key_len; -}; -struct predefined_key_element predefined_keys[] = { -""") - output_file.write("\n".join(look_up_table)) - output_file.write("\n};\n") - output_file.flush() + write_output_file(output_file, arrays, look_up_table) if __name__ == '__main__': main()