scripts/generate_test_keys.py - third_party/github/Mbed-TLS/mbedtls-framework - Git at Google

 #!/usr/bin/env python3

 # Copyright The Mbed TLS Contributors
 # SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later

 """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, List, Tuple
 import re
 import argparse
 from mbedtls_framework.asymmetric_key_data import ASYMMETRIC_KEY_DATA
 from mbedtls_framework import build_tree

 BYTES_PER_LINE = 16

 def c_byte_array_literal_content(array_name: str, key_data: bytes) -> Iterator[str]:
     """Return C code that defines array_name as a byte array with the given content."""
     yield 'static const unsigned char '
     yield array_name
     yield '[] = {'
     for index in range(0, len(key_data), BYTES_PER_LINE):
         yield '\n   '
         for b in key_data[index:index + BYTES_PER_LINE]:
             yield ' {:#04x},'.format(b)
     yield '\n};'

 def convert_der_to_c(array_name: str, key_data: bytes) -> str:
     return ''.join(c_byte_array_literal_content(array_name, key_data))

 def get_key_type(key_type: str) -> str:
     """Short name for a PSA key type."""
     if key_type.startswith('PSA_KEY_TYPE_ECC_'):
         return "ec"
     elif key_type.startswith('PSA_KEY_TYPE_ML_DSA_'):
         return "mldsa"
     elif key_type.startswith('PSA_KEY_TYPE_ML_KEM_'):
         return "mlkem"
     elif key_type.startswith('PSA_KEY_TYPE_RSA_'):
         return "rsa"
     elif key_type.startswith('PSA_KEY_TYPE_SLH_DSA_'):
         return "slhdsa"
     else:
         raise Exception(f"Unhandled key type {key_type}")

 def get_ec_key_family(key: str) -> str:
     """Extract "PSA_ECC_xxx" from "PSA_KEY_TYPE_ECC_ttt(PSA_ECC_xxx)"."""
     match = re.search(r'.*\((.*)\)', key)
     if match is None:
         raise Exception("Unable to get EC family from {}".format(key))
     return match.group(1)

 # Legacy EC group ID do not support all the key types that PSA does, so the
 # following dictionaries are used for:
 # - getting prefix/suffix for legacy curve names
 # - understand if the curve is supported in legacy symbols (MBEDTLS_ECP_DP_...)
 EC_NAME_CONVERSION = {
     'PSA_ECC_FAMILY_SECP_K1': {
         192: ('secp', 'k1'),
         256: ('secp', 'k1')
     },
     'PSA_ECC_FAMILY_SECP_R1': {
         192: ('secp', 'r1'),
         224: ('secp', 'r1'),
         256: ('secp', 'r1'),
         384: ('secp', 'r1'),
         521: ('secp', 'r1')
     },
     'PSA_ECC_FAMILY_BRAINPOOL_P_R1': {
         256: ('bp', 'r1'),
         384: ('bp', 'r1'),
         512: ('bp', 'r1')
     },
     'PSA_ECC_FAMILY_MONTGOMERY': {
         255: ('curve', '19'),
         448: ('curve', '')
     }
 }

 def get_ec_curve_name(priv_key: str, bits: int) -> str:
     """Short name for an elliptic curve key type."""
     ec_family = get_ec_key_family(priv_key)
     try:
         prefix = EC_NAME_CONVERSION[ec_family][bits][0]
         suffix = EC_NAME_CONVERSION[ec_family][bits][1]
     except KeyError:
         return ""
     return prefix + str(bits) + suffix

 def get_slh_dsa_family(key_type: str) -> str:
     """Short name from an SLH-DSA family."""
     m = re.search(r'PSA_SLH_FAMILY_(\w+)', key_type)
     assert m
     return m.group(1).replace('_', '').lower()

 def get_look_up_table_entry(key_type: str, group_id_or_keybits: str,
                             priv_array_name: str, pub_array_name: str) -> Iterator[str]:
     """Yield C code lines for the definition of a key pair and its matching public key."""
     if key_type == "ec":
         yield "    {{ {}, 0,\n".format(group_id_or_keybits)
     else:
         yield "    {{ 0, {},\n".format(group_id_or_keybits)
     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: str, look_up_table: str):
     """Write generated content to the output file"""
     with open(output_file_name, 'wt') as output:
         output.write("""\
 /*********************************************************************************
  * This file was automatically generated from framework/scripts/generate_test_keys.py.
  * Please do not edit it manually.
  *********************************************************************************/

 #ifndef TEST_TEST_KEYS_H
 #define TEST_TEST_KEYS_H

 #if TF_PSA_CRYPTO_VERSION_MAJOR >= 1
 #include <tf_psa_crypto_common.h>
 #include <mbedtls/private/ecp.h>
 #else
 #include <common.h>
 #include <mbedtls/ecp.h>
 #endif

 """)
         output.write(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;
 }};

 MBEDTLS_MAYBE_UNUSED static struct predefined_key_element predefined_keys[] = {{
 {}
 }};

 #endif /* TEST_TEST_KEYS_H */

 /* End of generated file */
 """.format(look_up_table))

 def collect_keys() -> Tuple[str, str]:
     """"
     This function reads key data from ASYMMETRIC_KEY_DATA and, only for the
     keys supported in legacy ECP/RSA modules, it returns 2 strings:
     - the 1st contains C arrays declaration of these keys and
     - the 2nd contains the final look-up table for all these arrays.
     """
     arrays = []
     look_up_table = []

     # Get a list of private keys only in order to get a single item for every
     # (key type, key bits) pair. We know that ASYMMETRIC_KEY_DATA
     # contains also the public counterpart.
     priv_keys = [key for key in ASYMMETRIC_KEY_DATA if '_KEY_PAIR' in key]
     priv_keys = sorted(priv_keys)

     for priv_key in priv_keys:
         key_type = get_key_type(priv_key)
         pub_key = re.sub('_KEY_PAIR', '_PUBLIC_KEY', priv_key)

         for bits in ASYMMETRIC_KEY_DATA[priv_key]:
             if key_type == "ec":
                 curve = get_ec_curve_name(priv_key, bits)
                 # Ignore EC curves unsupported in legacy symbols
                 if curve == "":
                     continue
             # Create output array name
             if key_type == "ec":
                 array_name_base = "_".join(["test", key_type, curve])
             elif key_type == "slhdsa":
                 family = get_slh_dsa_family(priv_key)
                 array_name_base = "_".join(["test", key_type, family, str(bits)])
             else:
                 array_name_base = "_".join(["test", key_type, str(bits)])
             array_name_priv = array_name_base + "_priv"
             array_name_pub = array_name_base + "_pub"
             # Convert bytearray to C array
             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
             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()
             else:
                 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)))

     return ''.join(arrays), '\n'.join(look_up_table)

 def main() -> None:
     """Command line entry point."""
     default_output_path = build_tree.guess_project_root() + "/tests/include/test/test_keys.h"

     argparser = argparse.ArgumentParser()
     argparser.add_argument("--output", help="Output file", default=default_output_path)
     args = argparser.parse_args()

     output_file = args.output

     # Support for 224 bit EC curves (secp224r1 and secp224k1) was removed from
     # tf-psa-crypto. It only remains available for 3.6 LTS branch.
     if not build_tree.is_mbedtls_3_6():
         del EC_NAME_CONVERSION['PSA_ECC_FAMILY_SECP_R1'][224]
         del EC_NAME_CONVERSION['PSA_ECC_FAMILY_SECP_R1'][192]
         del EC_NAME_CONVERSION['PSA_ECC_FAMILY_SECP_K1'][192]

     arrays, look_up_table = collect_keys()

     write_output_file(output_file, arrays, look_up_table)

 if __name__ == '__main__':
     main()
	#!/usr/bin/env python3

	# Copyright The Mbed TLS Contributors
	# SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later

	"""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, List, Tuple
	import re
	import argparse
	from mbedtls_framework.asymmetric_key_data import ASYMMETRIC_KEY_DATA
	from mbedtls_framework import build_tree

	BYTES_PER_LINE = 16

	def c_byte_array_literal_content(array_name: str, key_data: bytes) -> Iterator[str]:
	"""Return C code that defines array_name as a byte array with the given content."""
	yield 'static const unsigned char '
	yield array_name
	yield '[] = {'
	for index in range(0, len(key_data), BYTES_PER_LINE):
	yield '\n '
	for b in key_data[index:index + BYTES_PER_LINE]:
	yield ' {:#04x},'.format(b)
	yield '\n};'

	def convert_der_to_c(array_name: str, key_data: bytes) -> str:
	return ''.join(c_byte_array_literal_content(array_name, key_data))

	def get_key_type(key_type: str) -> str:
	"""Short name for a PSA key type."""
	if key_type.startswith('PSA_KEY_TYPE_ECC_'):
	return "ec"
	elif key_type.startswith('PSA_KEY_TYPE_ML_DSA_'):
	return "mldsa"
	elif key_type.startswith('PSA_KEY_TYPE_ML_KEM_'):
	return "mlkem"
	elif key_type.startswith('PSA_KEY_TYPE_RSA_'):
	return "rsa"
	elif key_type.startswith('PSA_KEY_TYPE_SLH_DSA_'):
	return "slhdsa"
	else:
	raise Exception(f"Unhandled key type {key_type}")

	def get_ec_key_family(key: str) -> str:
	"""Extract "PSA_ECC_xxx" from "PSA_KEY_TYPE_ECC_ttt(PSA_ECC_xxx)"."""
	match = re.search(r'.\((.)\)', key)
	if match is None:
	raise Exception("Unable to get EC family from {}".format(key))
	return match.group(1)

	# Legacy EC group ID do not support all the key types that PSA does, so the
	# following dictionaries are used for:
	# - getting prefix/suffix for legacy curve names
	# - understand if the curve is supported in legacy symbols (MBEDTLS_ECP_DP_...)
	EC_NAME_CONVERSION = {
	'PSA_ECC_FAMILY_SECP_K1': {
	192: ('secp', 'k1'),
	256: ('secp', 'k1')
	},
	'PSA_ECC_FAMILY_SECP_R1': {
	192: ('secp', 'r1'),
	224: ('secp', 'r1'),
	256: ('secp', 'r1'),
	384: ('secp', 'r1'),
	521: ('secp', 'r1')
	},
	'PSA_ECC_FAMILY_BRAINPOOL_P_R1': {
	256: ('bp', 'r1'),
	384: ('bp', 'r1'),
	512: ('bp', 'r1')
	},
	'PSA_ECC_FAMILY_MONTGOMERY': {
	255: ('curve', '19'),
	448: ('curve', '')
	}
	}

	def get_ec_curve_name(priv_key: str, bits: int) -> str:
	"""Short name for an elliptic curve key type."""
	ec_family = get_ec_key_family(priv_key)
	try:
	prefix = EC_NAME_CONVERSION[ec_family][bits][0]
	suffix = EC_NAME_CONVERSION[ec_family][bits][1]
	except KeyError:
	return ""
	return prefix + str(bits) + suffix

	def get_slh_dsa_family(key_type: str) -> str:
	"""Short name from an SLH-DSA family."""
	m = re.search(r'PSA_SLH_FAMILY_(\w+)', key_type)
	assert m
	return m.group(1).replace('_', '').lower()

	def get_look_up_table_entry(key_type: str, group_id_or_keybits: str,
	priv_array_name: str, pub_array_name: str) -> Iterator[str]:
	"""Yield C code lines for the definition of a key pair and its matching public key."""
	if key_type == "ec":
	yield " {{ {}, 0,\n".format(group_id_or_keybits)
	else:
	yield " {{ 0, {},\n".format(group_id_or_keybits)
	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: str, look_up_table: str):
	"""Write generated content to the output file"""
	with open(output_file_name, 'wt') as output:
	output.write("""\
	/*********************************************************************************
	* This file was automatically generated from framework/scripts/generate_test_keys.py.
	* Please do not edit it manually.
	*********************************************************************************/

	#ifndef TEST_TEST_KEYS_H
	#define TEST_TEST_KEYS_H

	#if TF_PSA_CRYPTO_VERSION_MAJOR >= 1
	#include <tf_psa_crypto_common.h>
	#include <mbedtls/private/ecp.h>
	#else
	#include <common.h>
	#include <mbedtls/ecp.h>
	#endif

	""")
	output.write(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;
	}};

	MBEDTLS_MAYBE_UNUSED static struct predefined_key_element predefined_keys[] = {{
	{}
	}};

	#endif /* TEST_TEST_KEYS_H */

	/* End of generated file */
	""".format(look_up_table))

	def collect_keys() -> Tuple[str, str]:
	""""
	This function reads key data from ASYMMETRIC_KEY_DATA and, only for the
	keys supported in legacy ECP/RSA modules, it returns 2 strings:
	- the 1st contains C arrays declaration of these keys and
	- the 2nd contains the final look-up table for all these arrays.
	"""
	arrays = []
	look_up_table = []

	# Get a list of private keys only in order to get a single item for every
	# (key type, key bits) pair. We know that ASYMMETRIC_KEY_DATA
	# contains also the public counterpart.
	priv_keys = [key for key in ASYMMETRIC_KEY_DATA if '_KEY_PAIR' in key]
	priv_keys = sorted(priv_keys)

	for priv_key in priv_keys:
	key_type = get_key_type(priv_key)
	pub_key = re.sub('_KEY_PAIR', '_PUBLIC_KEY', priv_key)

	for bits in ASYMMETRIC_KEY_DATA[priv_key]:
	if key_type == "ec":
	curve = get_ec_curve_name(priv_key, bits)
	# Ignore EC curves unsupported in legacy symbols
	if curve == "":
	continue
	# Create output array name
	if key_type == "ec":
	array_name_base = "_".join(["test", key_type, curve])
	elif key_type == "slhdsa":
	family = get_slh_dsa_family(priv_key)
	array_name_base = "_".join(["test", key_type, family, str(bits)])
	else:
	array_name_base = "_".join(["test", key_type, str(bits)])
	array_name_priv = array_name_base + "_priv"
	array_name_pub = array_name_base + "_pub"
	# Convert bytearray to C array
	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
	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()
	else:
	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)))

	return ''.join(arrays), '\n'.join(look_up_table)

	def main() -> None:
	"""Command line entry point."""
	default_output_path = build_tree.guess_project_root() + "/tests/include/test/test_keys.h"

	argparser = argparse.ArgumentParser()
	argparser.add_argument("--output", help="Output file", default=default_output_path)
	args = argparser.parse_args()

	output_file = args.output

	# Support for 224 bit EC curves (secp224r1 and secp224k1) was removed from
	# tf-psa-crypto. It only remains available for 3.6 LTS branch.
	if not build_tree.is_mbedtls_3_6():
	del EC_NAME_CONVERSION['PSA_ECC_FAMILY_SECP_R1'][224]
	del EC_NAME_CONVERSION['PSA_ECC_FAMILY_SECP_R1'][192]
	del EC_NAME_CONVERSION['PSA_ECC_FAMILY_SECP_K1'][192]

	arrays, look_up_table = collect_keys()

	write_output_file(output_file, arrays, look_up_table)

	if __name__ == '__main__':
	main()