File: generate_test_keys.py

package info (click to toggle)
mbedtls 3.6.4-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 50,424 kB
  • sloc: ansic: 164,526; sh: 25,295; python: 14,825; makefile: 2,761; perl: 1,043; tcl: 4
file content (184 lines) | stat: -rwxr-xr-x 6,652 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
#!/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.build_tree import guess_project_root

BYTES_PER_LINE = 16

def c_byte_array_literal_content(array_name: str, key_data: bytes) -> Iterator[str]:
    yield '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: str) -> str:
    if re.match('PSA_KEY_TYPE_RSA_.*', key):
        return "rsa"
    elif re.match('PSA_KEY_TYPE_ECC_.*', key):
        return "ec"
    else:
        print("Unhandled key type {}".format(key))
        return "unknown"

def get_ec_key_family(key: str) -> str:
    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'),
        224: ('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:
    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_look_up_table_entry(key_type: str, group_id_or_keybits: str,
                            priv_array_name: str, pub_array_name: str) -> Iterator[str]:
    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):
    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.
 *********************************************************************************/
""")
        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;
}};

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

/* 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)
        # Ignore keys which are not EC or RSA
        if key_type == "unknown":
            continue

        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 == "rsa":
                array_name_base = "_".join(["test", key_type, str(bits)])
            else:
                array_name_base = "_".join(["test", key_type, curve])
            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:
    default_output_path = 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

    arrays, look_up_table = collect_keys()

    write_output_file(output_file, arrays, look_up_table)

if __name__ == '__main__':
    main()