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
11sig_tee_alg = {'TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256': 0x70414930,
12               'TEE_ALG_RSASSA_PKCS1_V1_5_SHA256': 0x70004830}
13
14enc_tee_alg = {'TEE_ALG_AES_GCM': 0x40000810}
15
16enc_key_type = {'SHDR_ENC_KEY_DEV_SPECIFIC': 0x0,
17                'SHDR_ENC_KEY_CLASS_WIDE': 0x1}
18
19TEE_ATTR_RSA_MODULUS = 0xD0000130
20TEE_ATTR_RSA_PUBLIC_EXPONENT = 0xD0000230
21
22SHDR_BOOTSTRAP_TA = 1
23SHDR_ENCRYPTED_TA = 2
24SHDR_SUBKEY = 3
25SHDR_MAGIC = 0x4f545348
26SHDR_SIZE = 20
27SK_HDR_SIZE = 20
28EHDR_SIZE = 12
29UUID_SIZE = 16
30# Use 12 bytes for nonce per recommendation
31NONCE_SIZE = 12
32TAG_SIZE = 16
33
34
35def value_to_key(db, val):
36    for k, v in db.items():
37        if v == val:
38            return k
39
40
41def uuid_v5_sha512(namespace_bytes, name):
42    from cryptography.hazmat.primitives import hashes
43    from uuid import UUID
44
45    h = hashes.Hash(hashes.SHA512())
46    h.update(namespace_bytes + bytes(name, 'utf-8'))
47    digest = h.finalize()
48    return UUID(bytes=digest[:16], version=5)
49
50
51def name_img_to_str(name_img):
52    return name_img.decode().split('\x00', 1)[0]
53
54
55def uuid_parse(s):
56    from uuid import UUID
57    return UUID(s)
58
59
60def int_parse(str):
61    return int(str, 0)
62
63
64def get_args():
65    import argparse
66    import textwrap
67
68    class OnlyOne(argparse.Action):
69        def __call__(self, parser, namespace, values, option_string=None):
70            a = self.dest + '_assigned'
71            if getattr(namespace, a, False):
72                raise argparse.ArgumentError(self, 'Can only be given once')
73            setattr(namespace, a, True)
74            setattr(namespace, self.dest, values)
75
76    def arg_add_uuid(parser):
77        parser.add_argument(
78            '--uuid', required=True, type=uuid_parse,
79            help='String UUID of the TA')
80
81    def arg_add_key(parser):
82        parser.add_argument(
83            '--key', required=True, help='''
84                Name of signing and verification key file (PEM format) or an
85                Amazon Resource Name (arn:) of an AWS KMS asymmetric key.
86                At least public key for the commands digest, stitch, and
87                verify, else a private key''')
88
89    def arg_add_enc_key(parser):
90        parser.add_argument(
91            '--enc-key', required=False, help='Encryption key string')
92
93    def arg_add_enc_key_type(parser):
94        parser.add_argument(
95            '--enc-key-type', required=False,
96            default='SHDR_ENC_KEY_DEV_SPECIFIC',
97            choices=list(enc_key_type.keys()), help='''
98                Encryption key type,
99                Defaults to SHDR_ENC_KEY_DEV_SPECIFIC.''')
100
101    def arg_add_ta_version(parser):
102        parser.add_argument(
103            '--ta-version', required=False, type=int_parse, default=0, help='''
104                TA version stored as a 32-bit unsigned integer and used for
105                rollback protection of TA install in the secure database.
106                Defaults to 0.''')
107
108    def arg_add_sig(parser):
109        parser.add_argument(
110            '--sig', required=True, dest='sigf',
111            help='Name of signature input file, defaults to <UUID>.sig')
112
113    def arg_add_dig(parser):
114        parser.add_argument(
115            '--dig', required=True, dest='digf',
116            help='Name of digest output file, defaults to <UUID>.dig')
117
118    def arg_add_in(parser):
119        parser.add_argument(
120            '--in', required=False, dest='inf', help='''
121                Name of application input file, defaults to
122                <UUID>.stripped.elf''')
123
124    def arg_add_out(parser):
125        parser.add_argument(
126            '--out', required=True, dest='outf',
127            help='Name of application output file, defaults to <UUID>.ta')
128
129    def arg_add_algo(parser):
130        parser.add_argument(
131            '--algo', required=False, choices=list(sig_tee_alg.keys()),
132            default='TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256', help='''
133                The hash and signature algorithm.
134                Defaults to TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256.''')
135
136    def arg_add_subkey(parser):
137        parser.add_argument(
138            '--subkey', action=OnlyOne, help='Name of subkey input file')
139
140    def arg_add_name(parser):
141        parser.add_argument('--name',
142                            help='Input name for subspace of a subkey')
143
144    def arg_add_subkey_uuid_in(parser):
145        parser.add_argument(
146            '--in', required=True, dest='inf',
147            help='Name of subkey input file')
148
149    def arg_add_max_depth(parser):
150        parser.add_argument(
151            '--max-depth', required=False, type=int_parse, help='''
152            Max depth of subkeys below this subkey''')
153
154    def arg_add_name_size(parser):
155        parser.add_argument(
156            '--name-size', required=True, type=int_parse, help='''
157            Size of (unsigned) input name for subspace of a subkey.
158            Set to 0 to create an identity subkey (a subkey having
159            the same UUID as the next subkey or TA)''')
160
161    def arg_add_subkey_version(parser):
162        parser.add_argument(
163            '--subkey-version', required=False, type=int_parse, default=0,
164            help='Subkey version used for rollback protection')
165
166    def arg_add_subkey_in(parser):
167        parser.add_argument(
168            '--in', required=True, dest='inf', help='''
169            Name of PEM file with the public key of the new subkey''')
170
171    def arg_add_subkey_out(parser):
172        parser.add_argument(
173            '--out', required=True, dest='outf',
174            help='Name of subkey output file')
175
176    def get_outf_default(parsed):
177        return str(parsed.uuid) + '.ta'
178
179    def get_inf_default(parsed):
180        return str(parsed.uuid) + '.stripped.elf'
181
182    def get_sigf_default(parsed):
183        return str(parsed.uuid) + '.sig'
184
185    def get_digf_default(parsed):
186        return str(parsed.uuid) + '.dig'
187
188    def assign_default_value(parsed, attr, func):
189        if hasattr(parsed, attr) and getattr(parsed, attr) is None:
190            setattr(parsed, attr, func(parsed))
191
192    parser = argparse.ArgumentParser(
193        description='Sign and encrypt (optional) a Trusted Application ' +
194        ' for OP-TEE.',
195        usage='%(prog)s <command> ...',
196        epilog='<command> -h for detailed help')
197    subparsers = parser.add_subparsers(
198            title='valid commands, with possible aliases in ()',
199            dest='command', metavar='')
200
201    parser_sign_enc = subparsers.add_parser(
202        'sign-enc', prog=parser.prog + ' sign-enc',
203        help='Generate signed and optionally encrypted loadable TA image file')
204    parser_sign_enc.set_defaults(func=command_sign_enc)
205    arg_add_uuid(parser_sign_enc)
206    arg_add_ta_version(parser_sign_enc)
207    arg_add_in(parser_sign_enc)
208    arg_add_out(parser_sign_enc)
209    arg_add_key(parser_sign_enc)
210    arg_add_subkey(parser_sign_enc)
211    arg_add_name(parser_sign_enc)
212    arg_add_enc_key(parser_sign_enc)
213    arg_add_enc_key_type(parser_sign_enc)
214    arg_add_algo(parser_sign_enc)
215
216    parser_digest = subparsers.add_parser(
217        'digest', aliases=['generate-digest'], prog=parser.prog + ' digest',
218        formatter_class=argparse.RawDescriptionHelpFormatter,
219        help='Generate loadable TA binary image digest for offline signing',
220        epilog=textwrap.dedent('''\
221            example offline signing command using OpenSSL for algorithm
222            TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256:
223              base64 -d <UUID>.dig | \\
224              openssl pkeyutl -sign -inkey <KEYFILE>.pem \\
225                  -pkeyopt digest:sha256 -pkeyopt rsa_padding_mode:pss \\
226                  -pkeyopt rsa_pss_saltlen:digest \\
227                  -pkeyopt rsa_mgf1_md:sha256 | \\
228              base64 > <UUID>.sig
229
230            example offline signing command using OpenSSL for algorithm
231            TEE_ALG_RSASSA_PKCS1_V1_5_SHA256:
232              base64 -d <UUID>.dig | \\
233              openssl pkeyutl -sign -inkey <KEYFILE>.pem \\
234                  -pkeyopt digest:sha256 -pkeyopt rsa_padding_mode:pkcs1 | \\
235              base64 > <UUID>.sig
236            '''))
237    parser_digest.set_defaults(func=command_digest)
238    arg_add_uuid(parser_digest)
239    arg_add_ta_version(parser_digest)
240    arg_add_in(parser_digest)
241    arg_add_key(parser_digest)
242    arg_add_enc_key(parser_digest)
243    arg_add_enc_key_type(parser_digest)
244    arg_add_algo(parser_digest)
245    arg_add_dig(parser_digest)
246
247    parser_stitch = subparsers.add_parser(
248        'stitch', aliases=['stitch-ta'], prog=parser.prog + ' stich',
249        help='Generate loadable signed and encrypted TA binary image file' +
250        ' from TA raw image and its signature')
251    parser_stitch.set_defaults(func=command_stitch)
252    arg_add_uuid(parser_stitch)
253    arg_add_ta_version(parser_stitch)
254    arg_add_in(parser_stitch)
255    arg_add_key(parser_stitch)
256    arg_add_out(parser_stitch)
257    arg_add_enc_key(parser_stitch)
258    arg_add_enc_key_type(parser_stitch)
259    arg_add_algo(parser_stitch)
260    arg_add_sig(parser_stitch)
261
262    parser_verify = subparsers.add_parser(
263        'verify', prog=parser.prog + ' verify',
264        help='Verify signed TA binary')
265    parser_verify.set_defaults(func=command_verify)
266    arg_add_uuid(parser_verify)
267    arg_add_in(parser_verify)
268    arg_add_key(parser_verify)
269
270    parser_display = subparsers.add_parser(
271        'display', prog=parser.prog + ' display',
272        help='Parses and displays a signed TA binary')
273    parser_display.set_defaults(func=command_display)
274    arg_add_in(parser_display)
275
276    parser_subkey_uuid = subparsers.add_parser(
277        'subkey-uuid', prog=parser.prog + ' subkey-uuid',
278        help='calculate the UUID of next TA or subkey')
279    parser_subkey_uuid.set_defaults(func=command_subkey_uuid)
280    arg_add_subkey_uuid_in(parser_subkey_uuid)
281    arg_add_name(parser_subkey_uuid)
282
283    parser_sign_subkey = subparsers.add_parser(
284        'sign-subkey', prog=parser.prog + ' sign-subkey',
285        help='Sign a subkey')
286    parser_sign_subkey.set_defaults(func=command_sign_subkey)
287    arg_add_name(parser_sign_subkey)
288    arg_add_subkey_in(parser_sign_subkey)
289    arg_add_uuid(parser_sign_subkey)
290    arg_add_key(parser_sign_subkey)
291    arg_add_subkey_out(parser_sign_subkey)
292    arg_add_max_depth(parser_sign_subkey)
293    arg_add_name_size(parser_sign_subkey)
294    arg_add_subkey(parser_sign_subkey)
295    arg_add_subkey_version(parser_sign_subkey)
296    arg_add_algo(parser_sign_subkey)
297
298    argv = sys.argv[1:]
299    if (len(argv) > 0 and argv[0][0] == '-' and
300            argv[0] != '-h' and argv[0] != '--help'):
301        # The default sub-command is 'sign-enc' so add it to the parser
302        # if one is missing
303        argv = ['sign-enc'] + argv
304
305    parsed = parser.parse_args(argv)
306
307    if parsed.command is None:
308        parser.print_help()
309        sys.exit(1)
310
311    # Set a few defaults if defined for the current command
312    assign_default_value(parsed, 'inf', get_inf_default)
313    assign_default_value(parsed, 'outf', get_outf_default)
314    assign_default_value(parsed, 'sigf', get_sigf_default)
315    assign_default_value(parsed, 'digf', get_digf_default)
316
317    return parsed
318
319
320def load_asymmetric_key_img(data):
321    from cryptography.hazmat.backends import default_backend
322    from cryptography.hazmat.primitives.serialization import (
323        load_pem_private_key, load_pem_public_key)
324
325    try:
326        return load_pem_private_key(data, password=None,
327                                    backend=default_backend())
328    except ValueError:
329        return load_pem_public_key(data, backend=default_backend())
330
331
332def load_asymmetric_key(arg_key):
333    if arg_key.startswith('arn:'):
334        from sign_helper_kms import _RSAPrivateKeyInKMS
335        return _RSAPrivateKeyInKMS(arg_key)
336    else:
337        with open(arg_key, 'rb') as f:
338            return load_asymmetric_key_img(f.read())
339
340
341class BinaryImage:
342    def __init__(self, arg_inf, arg_key):
343        from cryptography.hazmat.primitives import hashes
344
345        # Exactly what inf is holding isn't determined a this stage
346        if isinstance(arg_inf, str):
347            with open(arg_inf, 'rb') as f:
348                self.inf = f.read()
349        else:
350            self.inf = arg_inf
351
352        if arg_key is None:
353            self.key = None
354        else:
355            if isinstance(arg_key, str):
356                self.key = load_asymmetric_key(arg_key)
357            else:
358                self.key = arg_key
359            self.sig_size = math.ceil(self.key.key_size / 8)
360
361        self.chosen_hash = hashes.SHA256()
362        self.hash_size = self.chosen_hash.digest_size
363
364    def __pack_img(self, img_type, sign_algo):
365        import struct
366
367        self.sig_algo = sign_algo
368        self.img_type = img_type
369        self.shdr = struct.pack('<IIIIHH', SHDR_MAGIC, img_type, len(self.img),
370                                sig_tee_alg[sign_algo], self.hash_size,
371                                self.sig_size)
372
373    def __calc_digest(self):
374        from cryptography.hazmat.backends import default_backend
375        from cryptography.hazmat.primitives import hashes
376
377        h = hashes.Hash(self.chosen_hash, default_backend())
378        h.update(self.shdr)
379        if hasattr(self, 'ta_uuid'):
380            h.update(self.ta_uuid)
381            h.update(self.ta_version)
382        if hasattr(self, 'ehdr'):
383            h.update(self.ehdr)
384            h.update(self.nonce)
385            h.update(self.tag)
386        h.update(self.img)
387        return h.finalize()
388
389    def encrypt_ta(self, enc_key, key_type, sig_algo, uuid, ta_version):
390        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
391        import struct
392        import os
393
394        self.img = self.inf
395
396        cipher = AESGCM(bytes.fromhex(enc_key))
397        self.nonce = os.urandom(NONCE_SIZE)
398        out = cipher.encrypt(self.nonce, self.img, None)
399        self.ciphertext = out[:-TAG_SIZE]
400        # Authentication Tag is always the last bytes
401        self.tag = out[-TAG_SIZE:]
402
403        enc_algo = enc_tee_alg['TEE_ALG_AES_GCM']
404        flags = enc_key_type[key_type]
405        self.ehdr = struct.pack('<IIHH', enc_algo, flags, len(self.nonce),
406                                len(self.tag))
407
408        self.__pack_img(SHDR_ENCRYPTED_TA, sig_algo)
409        self.ta_uuid = uuid.bytes
410        self.ta_version = struct.pack('<I', ta_version)
411        self.img_digest = self.__calc_digest()
412
413    def set_bootstrap_ta(self, sig_algo, uuid, ta_version):
414        import struct
415
416        self.img = self.inf
417        self.__pack_img(SHDR_BOOTSTRAP_TA, sig_algo)
418        self.ta_uuid = uuid.bytes
419        self.ta_version = struct.pack('<I', ta_version)
420        self.img_digest = self.__calc_digest()
421
422    def set_subkey(self, sign_algo, name, uuid, subkey_version, max_depth,
423                   name_size):
424        from cryptography.hazmat.primitives.asymmetric import rsa
425        import struct
426
427        self.subkey_name = name
428
429        subkey_key = load_asymmetric_key_img(self.inf)
430        if isinstance(subkey_key, rsa.RSAPrivateKey):
431            subkey_pkey = subkey_key.public_key()
432        else:
433            subkey_pkey = subkey_key
434
435        if max_depth is None:
436            if hasattr(self, 'previous_max_depth'):
437                if self.previous_max_depth <= 0:
438                    logger.error('Max depth of previous subkey is {}, '
439                                 .format(self.previous_max_depth) +
440                                 'cannot use a smaller value')
441                    sys.exit(1)
442
443                max_depth = self.previous_max_depth - 1
444            else:
445                max_depth = 0
446        else:
447            if (hasattr(self, 'previous_max_depth') and
448                    max_depth >= getattr(self, 'previous_max_depth')):
449                logger.error('Max depth of previous subkey is {} '
450                             .format(self.previous_max_depth) +
451                             'and the next value must be smaller')
452                sys.exit(1)
453
454        def int_to_bytes(x: int) -> bytes:
455            return x.to_bytes((x.bit_length() + 8) // 8, 'big')
456
457        n_bytes = int_to_bytes(subkey_pkey.public_numbers().n)
458        e_bytes = int_to_bytes(subkey_pkey.public_numbers().e)
459        attrs_end_offs = 16 + 5 * 4 + 2 * 3 * 4
460        shdr_subkey = struct.pack('<IIIIIIIIIII',
461                                  name_size, subkey_version,
462                                  max_depth, sig_tee_alg[sign_algo], 2,
463                                  TEE_ATTR_RSA_MODULUS,
464                                  attrs_end_offs, len(n_bytes),
465                                  TEE_ATTR_RSA_PUBLIC_EXPONENT,
466                                  attrs_end_offs + len(n_bytes),
467                                  len(e_bytes))
468        self.img = uuid.bytes + shdr_subkey + n_bytes + e_bytes
469        self.__pack_img(SHDR_SUBKEY, sign_algo)
470        self.img_digest = self.__calc_digest()
471
472    def parse(self):
473        from cryptography.hazmat.primitives.asymmetric import rsa
474        import struct
475
476        offs = 0
477        self.shdr = self.inf[offs:offs + SHDR_SIZE]
478        [magic, img_type, img_size, algo_value, hash_size,
479         sig_size] = struct.unpack('<IIIIHH', self.shdr)
480        offs += SHDR_SIZE
481
482        if magic != SHDR_MAGIC:
483            raise Exception("Unexpected magic: 0x{:08x}".format(magic))
484
485        if algo_value not in sig_tee_alg.values():
486            raise Exception('Unrecognized algorithm: 0x{:08x}'
487                            .format(algo_value))
488        self.sig_algo = value_to_key(sig_tee_alg, algo_value)
489
490        if hash_size != self.hash_size:
491            raise Exception("Unexpected digest len: {}".format(hash_size))
492
493        self.img_digest = self.inf[offs:offs + hash_size]
494        offs += hash_size
495        self.sig = self.inf[offs:offs + sig_size]
496        offs += sig_size
497
498        if img_type == SHDR_BOOTSTRAP_TA or img_type == SHDR_ENCRYPTED_TA:
499            self.ta_uuid = self.inf[offs:offs + UUID_SIZE]
500            offs += UUID_SIZE
501            self.ta_version = self.inf[offs:offs + 4]
502            offs += 4
503            if img_type == SHDR_ENCRYPTED_TA:
504                self.ehdr = self.inf[offs: offs + EHDR_SIZE]
505                offs += EHDR_SIZE
506                [enc_algo, flags, nonce_len,
507                 tag_len] = struct.unpack('<IIHH', self.ehdr)
508                if enc_value not in enc_tee_alg.values():
509                    raise Exception('Unrecognized encrypt algorithm: 0x{:08x}'
510                                    .format(enc_value))
511                if nonce_len != 12:
512                    raise Exception("Unexpected nonce len: {}"
513                                    .format(nonce_len))
514                self.nonce = self.inf[offs:offs + nonce_len]
515                offs += nonce_len
516
517                if tag_len != 16:
518                    raise Exception("Unexpected tag len: {}".format(tag_len))
519                self.tag = self.inf[-tag_len:]
520                self.ciphertext = self.inf[offs:-tag_len]
521                if len(self.ciphertext) != img_size:
522                    raise Exception("Unexpected ciphertext size: ",
523                                    "got {}, expected {}"
524                                    .format(len(self.ciphertext), img_size))
525                self.img = self.ciphertext
526            else:
527                self.img = self.inf[offs:]
528                if len(self.img) != img_size:
529                    raise Exception("Unexpected img size: got {}, expected {}"
530                                    .format(len(self.img), img_size))
531        elif img_type == SHDR_SUBKEY:
532            subkey_offs = offs
533            self.uuid = self.inf[offs:offs + UUID_SIZE]
534            offs += UUID_SIZE
535            self.subkey_hdr = self.inf[offs:offs + SK_HDR_SIZE]
536            [self.name_size, self.subkey_version, self.max_depth, self.algo,
537             self.attr_count] = struct.unpack('<IIIII', self.subkey_hdr)
538            offs += len(self.subkey_hdr)
539            self.attr = self.inf[offs:offs + img_size -
540                                 UUID_SIZE - len(self.subkey_hdr)]
541            offs += len(self.attr)
542            self.name_img = self.inf[offs:offs + self.name_size]
543            offs += self.name_size
544            self.next_inf = self.inf[offs:]
545
546            def find_attr(attr):
547                if self.attr_count <= 0:
548                    return None
549                for n in range(self.attr_count):
550                    o = subkey_offs + UUID_SIZE + SK_HDR_SIZE + n * 12
551                    [attr_value, attr_offs,
552                     attr_len] = struct.unpack('<III', self.inf[o: o + 12])
553                    if attr_value == attr:
554                        o = subkey_offs + attr_offs
555                        return self.inf[o:o + attr_len]
556                return None
557
558            n_bytes = find_attr(TEE_ATTR_RSA_MODULUS)
559            e_bytes = find_attr(TEE_ATTR_RSA_PUBLIC_EXPONENT)
560            e = int.from_bytes(e_bytes, 'big')
561            n = int.from_bytes(n_bytes, 'big')
562            self.subkey_key = rsa.RSAPublicNumbers(e, n).public_key()
563
564            self.img = self.inf[subkey_offs:offs - self.name_size]
565            if len(self.img) != img_size:
566                raise Exception("Unexpected img size: got {}, expected {}"
567                                .format(len(self.img), img_size))
568        else:
569            raise Exception("Unsupported image type: {}".format(img_type))
570
571    def display(self):
572        import binascii
573        import struct
574        import uuid
575
576        def display_ta():
577            nonlocal offs
578            ta_uuid = self.inf[offs:offs + UUID_SIZE]
579            print(' struct shdr_bootstrap_ta')
580            print('  uuid:       {}'.format(uuid.UUID(bytes=ta_uuid)))
581            offs += UUID_SIZE
582            [ta_version] = struct.unpack('<I', self.inf[offs:offs + 4])
583            print('  ta_version: {}'.format(ta_version))
584
585            offs += 4
586            if img_type == SHDR_ENCRYPTED_TA:
587                ehdr = self.inf[offs: offs + EHDR_SIZE]
588                offs += EHDR_SIZE
589                [enc_algo, flags, nonce_len,
590                 tag_len] = struct.unpack('<IIHH', ehdr)
591
592                print(' struct shdr_encrypted_ta')
593                enc_algo_name = 'Unkown'
594                if enc_algo in enc_tee_alg.values():
595                    enc_algo_name = value_to_key(enc_tee_alg, enc_algo)
596                print('  enc_algo:   0x{:08x} ({})'
597                      .format(enc_algo, enc_algo_name))
598
599                if enc_algo not in enc_tee_alg.values():
600                    raise Exception('Unrecognized encrypt algorithm: 0x{:08x}'
601                                    .format(enc_value))
602
603                flags_name = 'Unkown'
604                if flags in enc_key_type.values():
605                    flags_name = value_to_key(enc_key_type, flags)
606                print('  flags:      0x{:x} ({})'.format(flags, flags_name))
607
608                print('  iv_size:    {} (bytes)'.format(nonce_len))
609                if nonce_len != NONCE_SIZE:
610                    raise Exception("Unexpected nonce len: {}"
611                                    .format(nonce_len))
612                nonce = self.inf[offs:offs + nonce_len]
613                print('  iv:         {}'
614                      .format(binascii.hexlify(nonce).decode('ascii')))
615                offs += nonce_len
616
617                print('  tag_size:   {} (bytes)'.format(tag_len))
618                if tag_len != TAG_SIZE:
619                    raise Exception("Unexpected tag len: {}".format(tag_len))
620                tag = self.inf[-tag_len:]
621                print('  tag:        {}'
622                      .format(binascii.hexlify(tag).decode('ascii')))
623                ciphertext = self.inf[offs:-tag_len]
624                print(' TA offset:  {} (0x{:x}) bytes'.format(offs, offs))
625                print(' TA size:    {} (0x{:x}) bytes'
626                      .format(len(ciphertext), len(ciphertext)))
627                if len(ciphertext) != img_size:
628                    raise Exception("Unexpected ciphertext size: ",
629                                    "got {}, expected {}"
630                                    .format(len(ciphertext), img_size))
631                offs += tag_len
632            else:
633                img = self.inf[offs:]
634                print(' TA offset:  {} (0x{:x}) bytes'.format(offs, offs))
635                print(' TA size:    {} (0x{:x}) bytes'
636                      .format(len(img), len(img)))
637                if len(img) != img_size:
638                    raise Exception("Unexpected img size: got {}, expected {}"
639                                    .format(len(img), img_size))
640            offs += img_size
641
642        offs = 0
643        while offs < len(self.inf):
644            if offs > 0:
645                # name_size is the previous subkey header
646                name_img = self.inf[offs:offs + name_size]
647                print('  next name:  "{}"'.format(name_img_to_str(name_img)))
648                offs += name_size
649                print('Next header at offset: {} (0x{:x})'
650                      .format(offs, offs))
651
652            shdr = self.inf[offs:offs + SHDR_SIZE]
653            [magic, img_type, img_size, algo_value, hash_size,
654             sig_size] = struct.unpack('<IIIIHH', shdr)
655            offs += SHDR_SIZE
656
657            if magic != SHDR_MAGIC:
658                Exception("Unexpected magic: 0x{:08x}".format(magic))
659
660            img_type_name = 'Unknown'
661            if img_type == SHDR_BOOTSTRAP_TA:
662                print('Bootstrap TA')
663                img_type_name = 'SHDR_BOOTSTRAP_TA'
664            if img_type == SHDR_ENCRYPTED_TA:
665                print('Encrypted TA')
666                img_type_name = 'SHDR_ENCRYPTED_TA'
667            if img_type == SHDR_SUBKEY:
668                print('Subkey')
669                img_type_name = 'SHDR_SUBKEY'
670
671            algo_name = 'Unknown'
672            if algo_value in sig_tee_alg.values():
673                algo_name = value_to_key(sig_tee_alg, algo_value)
674
675            print(' struct shdr')
676            print('  magic:      0x{:08x}'.format(magic))
677            print('  img_type:   {} ({})'.format(img_type, img_type_name))
678            print('  img_size:   {} bytes'.format(img_size))
679            print('  algo:       0x{:08x} ({})'.format(algo_value, algo_name))
680            print('  hash_size:  {} bytes'.format(hash_size))
681            print('  sig_size:   {} bytes'.format(sig_size))
682
683            if algo_value not in sig_tee_alg.values():
684                raise Exception('Unrecognized algorithm: 0x{:08x}'
685                                .format(algo_value))
686
687            if hash_size != self.hash_size:
688                raise Exception("Unexpected digest len: {}".format(hash_size))
689
690            img_digest = self.inf[offs:offs + hash_size]
691            print('  hash:       {}'
692                  .format(binascii.hexlify(img_digest).decode('ascii')))
693            offs += hash_size
694            sig = self.inf[offs:offs + sig_size]
695            offs += sig_size
696
697            if img_type == SHDR_BOOTSTRAP_TA or img_type == SHDR_ENCRYPTED_TA:
698                display_ta()
699            elif img_type == SHDR_SUBKEY:
700                img_uuid = self.inf[offs:offs + UUID_SIZE]
701                img_subkey = self.inf[offs + UUID_SIZE:
702                                      offs + UUID_SIZE + SK_HDR_SIZE]
703                [name_size, subkey_version, max_depth, algo,
704                 attr_count] = struct.unpack('<IIIII', img_subkey)
705                if algo not in sig_tee_alg.values():
706                    raise Exception('Unrecognized algorithm: 0x{:08x}'
707                                    .format(algo))
708                algo_name = value_to_key(sig_tee_alg, algo)
709                print(' struct shdr_subkey')
710                print('  uuid:       {}'.format(uuid.UUID(bytes=img_uuid)))
711                print('  name_size:  {}'.format(name_size))
712                print('  subkey_version: {}'.format(subkey_version))
713                print('  max_depth:  {}'.format(max_depth))
714                print('  algo:       0x{:08x} ({})'.format(algo, algo_name))
715                print('  attr_count: {}'.format(attr_count))
716                offs += img_size
717            else:
718                raise Exception("Unsupported image type: {}".format(img_type))
719
720    def decrypt_ta(enc_key):
721        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
722
723        cipher = AESGCM(bytes.fromhex(enc_key))
724        self.img = cipher.decrypt(self.nonce, self.ciphertext, None)
725
726    def __get_padding(self):
727        from cryptography.hazmat.primitives.asymmetric import padding
728
729        if self.sig_algo == 'TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256':
730            pad = padding.PSS(mgf=padding.MGF1(self.chosen_hash),
731                              salt_length=self.hash_size)
732        elif self.sig_algo == 'TEE_ALG_RSASSA_PKCS1_V1_5_SHA256':
733            pad = padding.PKCS1v15()
734
735        return pad
736
737    def sign(self):
738        from cryptography.hazmat.primitives.asymmetric import utils
739        from cryptography.hazmat.primitives.asymmetric import rsa
740
741        if not isinstance(self.key, rsa.RSAPrivateKey):
742            logger.error('Provided key cannot be used for signing, ' +
743                         'please use offline-signing mode.')
744            sys.exit(1)
745        else:
746            self.sig = self.key.sign(self.img_digest, self.__get_padding(),
747                                     utils.Prehashed(self.chosen_hash))
748
749            if len(self.sig) != self.sig_size:
750                raise Exception(("Actual signature length is not equal to ",
751                                 "the computed one: {} != {}").
752                                format(len(self.sig), self.sig_size))
753
754    def add_signature(self, sigf):
755        import base64
756
757        with open(sigf, 'r') as f:
758            self.sig = base64.b64decode(f.read())
759
760        if len(self.sig) != self.sig_size:
761            raise Exception(("Actual signature length is not equal to ",
762                             "the expected one: {} != {}").
763                            format(len(self.sig), self.sig_size))
764
765    def verify_signature(self):
766        from cryptography.hazmat.primitives.asymmetric import utils
767        from cryptography.hazmat.primitives.asymmetric import rsa
768        from cryptography import exceptions
769
770        if isinstance(self.key, rsa.RSAPrivateKey):
771            pkey = self.key.public_key()
772        else:
773            pkey = self.key
774
775        try:
776            pkey.verify(self.sig, self.img_digest, self.__get_padding(),
777                        utils.Prehashed(self.chosen_hash))
778        except exceptions.InvalidSignature:
779            logger.error('Verification failed, ignoring given signature.')
780            sys.exit(1)
781
782    def verify_digest(self):
783        if self.img_digest != self.__calc_digest():
784            raise Exception('Hash digest does not match')
785
786    def verify_uuid(self, uuid):
787        if self.ta_uuid != uuid.bytes:
788            raise Exception('UUID does not match')
789
790    def add_subkey(self, subkey_file, name):
791        sk_image = BinaryImage(subkey_file, None)
792        self.subkey_img = sk_image.inf
793        sk_image.parse()
794        if not hasattr(sk_image, 'next_inf'):
795            logger.error('Invalid subkey file')
796            sys.exit(1)
797        while len(sk_image.next_inf) > 0:
798            sk_image = BinaryImage(sk_image.next_inf, None)
799            sk_image.parse()
800
801        if name is None:
802            name = ''
803        self.previous_max_depth = sk_image.max_depth
804        self.name_img = str.encode(name).ljust(sk_image.name_size, b'\0')
805
806    def write(self, outf):
807        with open(outf, 'wb') as f:
808            if hasattr(self, 'subkey_img'):
809                f.write(self.subkey_img)
810                f.write(self.name_img)
811            f.write(self.shdr)
812            f.write(self.img_digest)
813            f.write(self.sig)
814            if hasattr(self, 'ta_uuid'):
815                f.write(self.ta_uuid)
816                f.write(self.ta_version)
817            if hasattr(self, 'ehdr'):
818                f.write(self.ehdr)
819                f.write(self.nonce)
820                f.write(self.tag)
821                f.write(self.ciphertext)
822            else:
823                f.write(self.img)
824
825
826def load_ta_image(args):
827    ta_image = BinaryImage(args.inf, args.key)
828
829    if args.enc_key:
830        ta_image.encrypt_ta(args.enc_key, args.enc_key_type,
831                            args.algo, args.uuid, args.ta_version)
832    else:
833        ta_image.set_bootstrap_ta(args.algo, args.uuid, args.ta_version)
834
835    return ta_image
836
837
838def command_sign_enc(args):
839    ta_image = load_ta_image(args)
840    if args.subkey:
841        ta_image.add_subkey(args.subkey, args.name)
842    ta_image.sign()
843    ta_image.write(args.outf)
844    logger.info('Successfully signed application.')
845
846
847def command_sign_subkey(args):
848    image = BinaryImage(args.inf, args.key)
849    if args.subkey:
850        image.add_subkey(args.subkey, args.name)
851    image.set_subkey(args.algo, args.name, args.uuid, args.subkey_version,
852                     args.max_depth, args.name_size)
853    image.sign()
854    image.write(args.outf)
855    logger.info('Successfully signed subkey.')
856
857
858def command_digest(args):
859    import base64
860
861    ta_image = load_ta_image(args)
862    with open(args.digf, 'wb+') as digfile:
863        digfile.write(base64.b64encode(ta_image.img_digest))
864
865
866def command_stitch(args):
867    ta_image = load_ta_image(args)
868    ta_image.add_signature(args.sigf)
869    ta_image.verify_signature()
870    ta_image.write(args.outf)
871    logger.info('Successfully applied signature.')
872
873
874def command_verify(args):
875    import uuid
876
877    image = BinaryImage(args.inf, args.key)
878    next_uuid = None
879    max_depth = -1
880    while True:
881        image.parse()
882        if hasattr(image, 'subkey_hdr'):  # Subkey
883            print('Subkey UUID: {}'.format(uuid.UUID(bytes=image.uuid)))
884            image.verify_signature()
885            image.verify_digest()
886            if next_uuid:
887                if uuid.UUID(bytes=image.uuid) != next_uuid:
888                    raise Exception('UUID {} does not match {}'
889                                    .format(uuid.UUID(bytes=image.uuid),
890                                            next_uuid))
891            if max_depth >= 0:
892                if image.max_depth < 0 or image.max_depth >= max_depth:
893                    raise Exception('Invalid max_depth {} not less than {}'
894                                    .format(image.max_depth, max_depth))
895            max_depth = image.max_depth
896            if len(image.next_inf) == 0:
897                logger.info('Subkey is correctly verified.')
898                return
899            if image.name_size > 0:
900                next_uuid = uuid_v5_sha512(image.uuid,
901                                           name_img_to_str(image.name_img))
902            else:
903                next_uuid = image.uuid
904            image = BinaryImage(image.next_inf, image.subkey_key)
905        else:  # TA
906            print('TA UUID: {}'.format(uuid.UUID(bytes=image.ta_uuid)))
907            if next_uuid:
908                if uuid.UUID(bytes=image.ta_uuid) != next_uuid:
909                    raise Exception('UUID {} does not match {}'
910                                    .format(uuid.UUID(bytes=image.ta_uuid),
911                                            next_uuid))
912            if hasattr(image, 'ciphertext'):
913                if args.enc_key is None:
914                    logger.error('--enc_key needed to decrypt TA')
915                    sys.exit(1)
916                image.decrypt_ta(args.enc_key)
917            image.verify_signature()
918            image.verify_digest()
919            image.verify_uuid(args.uuid)
920            logger.info('Trusted application is correctly verified.')
921            return
922
923
924def command_display(args):
925    ta_image = BinaryImage(args.inf, None)
926    ta_image.display()
927
928
929def command_subkey_uuid(args):
930    import uuid
931
932    sk_image = BinaryImage(args.inf, None)
933    sk_image.parse()
934    if not hasattr(sk_image, 'next_inf'):
935        logger.error('Invalid subkey file')
936        sys.exit(1)
937    print('Subkey UUID: {}'.format(uuid.UUID(bytes=sk_image.uuid)))
938    while len(sk_image.next_inf) > 0:
939        sk_image = BinaryImage(sk_image.next_inf, None)
940        sk_image.parse()
941        print('Subkey UUID: {}'.format(uuid.UUID(bytes=sk_image.uuid)))
942    if args.name:
943        if len(args.name) > sk_image.name_size:
944            logger.error('Length of name ({}) '.format(len(args.name)) +
945                         'is larger than max name size ({})'
946                         .format(sk_image.name_size))
947            sys.exit(1)
948        print('Next subkey UUID: {}'
949              .format(uuid_v5_sha512(sk_image.uuid, args.name)))
950    else:
951        print('Next subkey UUID unchanged: {}'
952              .format(uuid.UUID(bytes=sk_image.uuid)))
953
954
955def main():
956    import logging
957    import os
958
959    global logger
960    logging.basicConfig()
961    logger = logging.getLogger(os.path.basename(__file__))
962
963    args = get_args()
964    args.func(args)
965
966
967if __name__ == "__main__":
968    main()
969