1#!/usr/bin/env python3
2# SPDX-License-Identifier: BSD-2-Clause
3#
4# Copyright (c) 2015, 2017, 2019, Linaro Limited
5#
6
7import sys
8import math
9
10
11algo = {'TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256': 0x70414930,
12        'TEE_ALG_RSASSA_PKCS1_V1_5_SHA256': 0x70004830}
13
14enc_key_type = {'SHDR_ENC_KEY_DEV_SPECIFIC': 0x0,
15                'SHDR_ENC_KEY_CLASS_WIDE': 0x1}
16
17SHDR_BOOTSTRAP_TA = 1
18SHDR_ENCRYPTED_TA = 2
19SHDR_MAGIC = 0x4f545348
20SHDR_SIZE = 20
21
22
23def uuid_parse(s):
24    from uuid import UUID
25    return UUID(s)
26
27
28def int_parse(str):
29    return int(str, 0)
30
31
32def get_args(logger):
33    from argparse import ArgumentParser, RawDescriptionHelpFormatter
34    import textwrap
35    command_base = ['sign-enc', 'digest', 'stitch', 'verify']
36    command_aliases_digest = ['generate-digest']
37    command_aliases_stitch = ['stitch-ta']
38    command_aliases = command_aliases_digest + command_aliases_stitch
39    command_choices = command_base + command_aliases
40
41    dat = '[' + ', '.join(command_aliases_digest) + ']'
42    sat = '[' + ', '.join(command_aliases_stitch) + ']'
43
44    parser = ArgumentParser(
45        description='Sign and encrypt (optional) a Trusted Application for' +
46        ' OP-TEE.',
47        usage='\n   %(prog)s command [ arguments ]\n\n'
48
49        '   command:\n' +
50        '     sign-enc    Generate signed and optionally encrypted loadable' +
51        ' TA image file.\n' +
52        '                 Takes arguments --uuid, --ta-version, --in, --out,' +
53        ' --key,\n' +
54        '                 --enc-key (optional) and' +
55        ' --enc-key-type (optional).\n' +
56        '     digest      Generate loadable TA binary image digest' +
57        ' for offline\n' +
58        '                 signing. Takes arguments --uuid, --ta-version,' +
59        ' --in, --key,\n'
60        '                 --enc-key (optional), --enc-key-type (optional),' +
61        ' --algo (optional) and --dig.\n' +
62        '     stitch      Generate loadable signed and encrypted TA binary' +
63        ' image file from\n' +
64        '                 TA raw image and its signature. Takes' +
65        ' arguments --uuid, --in, --key, --out,\n' +
66        '                 --enc-key (optional), --enc-key-type (optional),\n' +
67        '                 --algo (optional) and --sig.\n' +
68        '     verify      Verify signed TA binary\n' +
69        '                 Takes arguments --uuid, --in, --key\n\n' +
70        '   %(prog)s --help  show available commands and arguments\n\n',
71        formatter_class=RawDescriptionHelpFormatter,
72        epilog=textwrap.dedent('''\
73            If no command is given, the script will default to "sign-enc".
74
75            command aliases:
76              The command \'digest\' can be aliased by ''' + dat + '''
77              The command \'stitch\' can be aliased by ''' + sat + '\n' + '''
78            example offline signing command using OpenSSL for algorithm
79            TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256:
80              base64 -d <UUID>.dig | \\
81              openssl pkeyutl -sign -inkey <KEYFILE>.pem \\
82                  -pkeyopt digest:sha256 -pkeyopt rsa_padding_mode:pss \\
83                  -pkeyopt rsa_pss_saltlen:digest \\
84                  -pkeyopt rsa_mgf1_md:sha256 | \\
85              base64 > <UUID>.sig\n
86            example offline signing command using OpenSSL for algorithm
87            TEE_ALG_RSASSA_PKCS1_V1_5_SHA256:
88              base64 -d <UUID>.dig | \\
89              openssl pkeyutl -sign -inkey <KEYFILE>.pem \\
90                  -pkeyopt digest:sha256 -pkeyopt rsa_padding_mode:pkcs1 | \\
91              base64 > <UUID>.sig
92            '''))
93
94    parser.add_argument(
95        'command', choices=command_choices, nargs='?',
96        default='sign-enc',
97        help='Command, one of [' + ', '.join(command_base) + ']')
98    parser.add_argument('--uuid', required=True,
99                        type=uuid_parse, help='String UUID of the TA')
100    parser.add_argument('--key', required=True,
101                        help='Name of signing key file (PEM format) or an ' +
102                             'Amazon Resource Name (arn:) of an AWS KMS ' +
103                             'asymmetric key')
104    parser.add_argument('--enc-key', required=False,
105                        help='Encryption key string')
106    parser.add_argument(
107        '--enc-key-type', required=False, default='SHDR_ENC_KEY_DEV_SPECIFIC',
108        choices=list(enc_key_type.keys()),
109        help='Encryption key type.\n' +
110        '(SHDR_ENC_KEY_DEV_SPECIFIC or SHDR_ENC_KEY_CLASS_WIDE).\n' +
111        'Defaults to SHDR_ENC_KEY_DEV_SPECIFIC.')
112    parser.add_argument(
113        '--ta-version', required=False, type=int_parse, default=0,
114        help='TA version stored as a 32-bit unsigned integer and used for\n' +
115        'rollback protection of TA install in the secure database.\n' +
116        'Defaults to 0.')
117    parser.add_argument(
118        '--sig', required=False, dest='sigf',
119        help='Name of signature input file, defaults to <UUID>.sig')
120    parser.add_argument(
121        '--dig', required=False, dest='digf',
122        help='Name of digest output file, defaults to <UUID>.dig')
123    parser.add_argument(
124        '--in', required=True, dest='inf',
125        help='Name of application input file, defaults to <UUID>.stripped.elf')
126    parser.add_argument(
127        '--out', required=False, dest='outf',
128        help='Name of application output file, defaults to <UUID>.ta')
129    parser.add_argument('--algo', required=False, choices=list(algo.keys()),
130                        default='TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256',
131                        help='The hash and signature algorithm, ' +
132                        'defaults to TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256. ' +
133                        'Allowed values are: ' +
134                        ', '.join(list(algo.keys())), metavar='')
135
136    parsed = parser.parse_args()
137
138    # Check parameter combinations
139
140    if parsed.digf is None and \
141       parsed.outf is not None and \
142       parsed.command in ['digest'] + command_aliases_digest:
143        logger.error('A digest was requested, but argument --out was given.' +
144                     '  Did you mean:\n  ' +
145                     parser.prog+' --dig ' + parsed.outf + ' ...')
146        sys.exit(1)
147
148    if parsed.digf is not None \
149       and parsed.outf is not None \
150       and parsed.command in ['digest'] + command_aliases_digest:
151        logger.warn('A digest was requested, but arguments --dig and ' +
152                    '--out were given.\n' +
153                    '  --out will be ignored.')
154
155    # Set defaults for optional arguments.
156
157    if parsed.sigf is None:
158        parsed.sigf = str(parsed.uuid)+'.sig'
159    if parsed.digf is None:
160        parsed.digf = str(parsed.uuid)+'.dig'
161    if parsed.inf is None:
162        parsed.inf = str(parsed.uuid)+'.stripped.elf'
163    if parsed.outf is None:
164        parsed.outf = str(parsed.uuid)+'.ta'
165
166    return parsed
167
168
169def main():
170    from cryptography import exceptions
171    from cryptography.hazmat.backends import default_backend
172    from cryptography.hazmat.primitives import serialization
173    from cryptography.hazmat.primitives import hashes
174    from cryptography.hazmat.primitives.asymmetric import padding
175    from cryptography.hazmat.primitives.asymmetric import rsa
176    from cryptography.hazmat.primitives.asymmetric import utils
177    import base64
178    import logging
179    import os
180    import struct
181
182    logging.basicConfig()
183    logger = logging.getLogger(os.path.basename(__file__))
184
185    args = get_args(logger)
186
187    if args.key.startswith('arn:'):
188        from sign_helper_kms import _RSAPrivateKeyInKMS
189        key = _RSAPrivateKeyInKMS(args.key)
190    else:
191        with open(args.key, 'rb') as f:
192            data = f.read()
193
194            try:
195                key = serialization.load_pem_private_key(
196                          data,
197                          password=None,
198                          backend=default_backend())
199            except ValueError:
200                key = serialization.load_pem_public_key(
201                          data,
202                          backend=default_backend())
203
204    with open(args.inf, 'rb') as f:
205        img = f.read()
206
207    chosen_hash = hashes.SHA256()
208    h = hashes.Hash(chosen_hash, default_backend())
209
210    digest_len = chosen_hash.digest_size
211    sig_len = math.ceil(key.key_size / 8)
212
213    img_size = len(img)
214
215    hdr_version = args.ta_version  # struct shdr_bootstrap_ta::ta_version
216
217    magic = SHDR_MAGIC
218    if args.enc_key:
219        img_type = SHDR_ENCRYPTED_TA
220    else:
221        img_type = SHDR_BOOTSTRAP_TA
222
223    shdr = struct.pack('<IIIIHH',
224                       magic, img_type, img_size, algo[args.algo],
225                       digest_len, sig_len)
226    shdr_uuid = args.uuid.bytes
227    shdr_version = struct.pack('<I', hdr_version)
228
229    if args.enc_key:
230        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
231        cipher = AESGCM(bytes.fromhex(args.enc_key))
232        # Use 12 bytes for nonce per recommendation
233        nonce = os.urandom(12)
234        out = cipher.encrypt(nonce, img, None)
235        ciphertext = out[:-16]
236        # Authentication Tag is always the last 16 bytes
237        tag = out[-16:]
238
239        enc_algo = 0x40000810      # TEE_ALG_AES_GCM
240        flags = enc_key_type[args.enc_key_type]
241        ehdr = struct.pack('<IIHH',
242                           enc_algo, flags, len(nonce), len(tag))
243
244    h.update(shdr)
245    h.update(shdr_uuid)
246    h.update(shdr_version)
247    if args.enc_key:
248        h.update(ehdr)
249        h.update(nonce)
250        h.update(tag)
251    h.update(img)
252    img_digest = h.finalize()
253
254    def write_image_with_signature(sig):
255        with open(args.outf, 'wb') as f:
256            f.write(shdr)
257            f.write(img_digest)
258            f.write(sig)
259            f.write(shdr_uuid)
260            f.write(shdr_version)
261            if args.enc_key:
262                f.write(ehdr)
263                f.write(nonce)
264                f.write(tag)
265                f.write(ciphertext)
266            else:
267                f.write(img)
268
269    def sign_encrypt_ta():
270        if not isinstance(key, rsa.RSAPrivateKey):
271            logger.error('Provided key cannot be used for signing, ' +
272                         'please use offline-signing mode.')
273            sys.exit(1)
274        else:
275            if args.algo == 'TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256':
276                sig = key.sign(
277                    img_digest,
278                    padding.PSS(
279                        mgf=padding.MGF1(chosen_hash),
280                        salt_length=digest_len
281                    ),
282                    utils.Prehashed(chosen_hash)
283                )
284            elif args.algo == 'TEE_ALG_RSASSA_PKCS1_V1_5_SHA256':
285                sig = key.sign(
286                    img_digest,
287                    padding.PKCS1v15(),
288                    utils.Prehashed(chosen_hash)
289                )
290
291            if len(sig) != sig_len:
292                raise Exception(("Actual signature length is not equal to ",
293                                 "the computed one: {} != {}").
294                                format(len(sig), sig_len))
295            write_image_with_signature(sig)
296            logger.info('Successfully signed application.')
297
298    def generate_digest():
299        with open(args.digf, 'wb+') as digfile:
300            digfile.write(base64.b64encode(img_digest))
301
302    def stitch_ta():
303        try:
304            with open(args.sigf, 'r') as sigfile:
305                sig = base64.b64decode(sigfile.read())
306        except IOError:
307            if not os.path.exists(args.digf):
308                generate_digest()
309            logger.error('No signature file found. Please sign\n %s\n' +
310                         'offline and place the signature at \n %s\n' +
311                         'or pass a different location ' +
312                         'using the --sig argument.\n',
313                         args.digf, args.sigf)
314            sys.exit(1)
315        else:
316            try:
317                if args.algo == 'TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256':
318                    key.verify(
319                        sig,
320                        img_digest,
321                        padding.PSS(
322                            mgf=padding.MGF1(chosen_hash),
323                            salt_length=digest_len
324                        ),
325                        utils.Prehashed(chosen_hash)
326                    )
327                elif args.algo == 'TEE_ALG_RSASSA_PKCS1_V1_5_SHA256':
328                    key.verify(
329                        sig,
330                        img_digest,
331                        padding.PKCS1v15(),
332                        utils.Prehashed(chosen_hash)
333                    )
334            except exceptions.InvalidSignature:
335                logger.error('Verification failed, ignoring given signature.')
336                sys.exit(1)
337
338            write_image_with_signature(sig)
339            logger.info('Successfully applied signature.')
340
341    def verify_ta():
342        # Extract header
343        [magic,
344         img_type,
345         img_size,
346         algo_value,
347         digest_len,
348         sig_len] = struct.unpack('<IIIIHH', img[:SHDR_SIZE])
349
350        # Extract digest and signature
351        start, end = SHDR_SIZE, SHDR_SIZE + digest_len
352        digest = img[start:end]
353
354        start, end = end, SHDR_SIZE + digest_len + sig_len
355        signature = img[start:end]
356
357        # Extract UUID and TA version
358        start, end = end, end + 16 + 4
359        [uuid, ta_version] = struct.unpack('<16sI', img[start:end])
360
361        if magic != SHDR_MAGIC:
362            raise Exception("Unexpected magic: 0x{:08x}".format(magic))
363
364        if img_type != SHDR_BOOTSTRAP_TA:
365            raise Exception("Unsupported image type: {}".format(img_type))
366
367        if algo_value not in algo.values():
368            raise Exception('Unrecognized algorithm: 0x{:08x}'
369                            .format(algo_value))
370
371        # Verify signature against hash digest
372        if algo_value == 0x70414930:
373            key.verify(
374                signature,
375                digest,
376                padding.PSS(
377                    mgf=padding.MGF1(chosen_hash),
378                    salt_length=digest_len
379                ),
380                utils.Prehashed(chosen_hash)
381            )
382        else:
383            key.verify(
384                signature,
385                digest,
386                padding.PKCS1v15(),
387                utils.Prehashed(chosen_hash)
388            )
389
390        h = hashes.Hash(chosen_hash, default_backend())
391
392        # sizeof(struct shdr)
393        h.update(img[:SHDR_SIZE])
394
395        # sizeof(struct shdr_bootstrap_ta)
396        h.update(img[start:end])
397
398        # raw image
399        start = end
400        end += img_size
401        h.update(img[start:end])
402
403        if digest != h.finalize():
404            raise Exception('Hash digest does not match')
405
406        logger.info('Trusted application is correctly verified.')
407
408    # dispatch command
409    {
410        'sign-enc': sign_encrypt_ta,
411        'digest': generate_digest,
412        'generate-digest': generate_digest,
413        'stitch': stitch_ta,
414        'stitch-ta': stitch_ta,
415        'verify': verify_ta,
416    }.get(args.command, 'sign_encrypt_ta')()
417
418
419if __name__ == "__main__":
420    main()
421