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 + ' stitch',
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    arg_add_enc_key(parser_verify)
270
271    parser_display = subparsers.add_parser(
272        'display', prog=parser.prog + ' display',
273        help='Parses and displays a signed TA binary')
274    parser_display.set_defaults(func=command_display)
275    arg_add_in(parser_display)
276
277    parser_subkey_uuid = subparsers.add_parser(
278        'subkey-uuid', prog=parser.prog + ' subkey-uuid',
279        help='calculate the UUID of next TA or subkey')
280    parser_subkey_uuid.set_defaults(func=command_subkey_uuid)
281    arg_add_subkey_uuid_in(parser_subkey_uuid)
282    arg_add_name(parser_subkey_uuid)
283
284    parser_sign_subkey = subparsers.add_parser(
285        'sign-subkey', prog=parser.prog + ' sign-subkey',
286        help='Sign a subkey')
287    parser_sign_subkey.set_defaults(func=command_sign_subkey)
288    arg_add_name(parser_sign_subkey)
289    arg_add_subkey_in(parser_sign_subkey)
290    arg_add_uuid(parser_sign_subkey)
291    arg_add_key(parser_sign_subkey)
292    arg_add_subkey_out(parser_sign_subkey)
293    arg_add_max_depth(parser_sign_subkey)
294    arg_add_name_size(parser_sign_subkey)
295    arg_add_subkey(parser_sign_subkey)
296    arg_add_subkey_version(parser_sign_subkey)
297    arg_add_algo(parser_sign_subkey)
298
299    argv = sys.argv[1:]
300    if (len(argv) > 0 and argv[0][0] == '-' and
301            argv[0] != '-h' and argv[0] != '--help'):
302        # The default sub-command is 'sign-enc' so add it to the parser
303        # if one is missing
304        argv = ['sign-enc'] + argv
305
306    parsed = parser.parse_args(argv)
307
308    if parsed.command is None:
309        parser.print_help()
310        sys.exit(1)
311
312    # Set a few defaults if defined for the current command
313    assign_default_value(parsed, 'inf', get_inf_default)
314    assign_default_value(parsed, 'outf', get_outf_default)
315    assign_default_value(parsed, 'sigf', get_sigf_default)
316    assign_default_value(parsed, 'digf', get_digf_default)
317
318    return parsed
319
320
321def load_asymmetric_key_img(data):
322    from cryptography.hazmat.backends import default_backend
323    from cryptography.hazmat.primitives.serialization import (
324        load_pem_private_key, load_pem_public_key)
325
326    try:
327        return load_pem_private_key(data, password=None,
328                                    backend=default_backend())
329    except ValueError:
330        return load_pem_public_key(data, backend=default_backend())
331
332
333def load_asymmetric_key(arg_key):
334    if arg_key.startswith('arn:'):
335        from sign_helper_kms import _RSAPrivateKeyInKMS
336        return _RSAPrivateKeyInKMS(arg_key)
337    else:
338        with open(arg_key, 'rb') as f:
339            return load_asymmetric_key_img(f.read())
340
341
342class BinaryImage:
343    def __init__(self, arg_inf, arg_key):
344        from cryptography.hazmat.primitives import hashes
345
346        # Exactly what inf is holding isn't determined a this stage
347        if isinstance(arg_inf, str):
348            with open(arg_inf, 'rb') as f:
349                self.inf = f.read()
350        else:
351            self.inf = arg_inf
352
353        if arg_key is None:
354            self.key = None
355        else:
356            if isinstance(arg_key, str):
357                self.key = load_asymmetric_key(arg_key)
358            else:
359                self.key = arg_key
360            self.sig_size = math.ceil(self.key.key_size / 8)
361
362        self.chosen_hash = hashes.SHA256()
363        self.hash_size = self.chosen_hash.digest_size
364
365    def __pack_img(self, img_type, sign_algo):
366        import struct
367
368        self.sig_algo = sign_algo
369        self.img_type = img_type
370        self.shdr = struct.pack('<IIIIHH', SHDR_MAGIC, img_type, len(self.img),
371                                sig_tee_alg[sign_algo], self.hash_size,
372                                self.sig_size)
373
374    def __calc_digest(self):
375        from cryptography.hazmat.backends import default_backend
376        from cryptography.hazmat.primitives import hashes
377
378        h = hashes.Hash(self.chosen_hash, default_backend())
379        h.update(self.shdr)
380        if hasattr(self, 'ta_uuid'):
381            h.update(self.ta_uuid)
382            h.update(self.ta_version)
383        if hasattr(self, 'ehdr'):
384            h.update(self.ehdr)
385            h.update(self.nonce)
386            h.update(self.tag)
387        h.update(self.img)
388        return h.finalize()
389
390    def encrypt_ta(self, enc_key, key_type, sig_algo, uuid, ta_version):
391        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
392        import struct
393        import os
394
395        self.img = self.inf
396
397        cipher = AESGCM(bytes.fromhex(enc_key))
398        self.nonce = os.urandom(NONCE_SIZE)
399        out = cipher.encrypt(self.nonce, self.img, None)
400        self.ciphertext = out[:-TAG_SIZE]
401        # Authentication Tag is always the last bytes
402        self.tag = out[-TAG_SIZE:]
403
404        enc_algo = enc_tee_alg['TEE_ALG_AES_GCM']
405        flags = enc_key_type[key_type]
406        self.ehdr = struct.pack('<IIHH', enc_algo, flags, len(self.nonce),
407                                len(self.tag))
408
409        self.__pack_img(SHDR_ENCRYPTED_TA, sig_algo)
410        self.ta_uuid = uuid.bytes
411        self.ta_version = struct.pack('<I', ta_version)
412        self.img_digest = self.__calc_digest()
413
414    def set_bootstrap_ta(self, sig_algo, uuid, ta_version):
415        import struct
416
417        self.img = self.inf
418        self.__pack_img(SHDR_BOOTSTRAP_TA, sig_algo)
419        self.ta_uuid = uuid.bytes
420        self.ta_version = struct.pack('<I', ta_version)
421        self.img_digest = self.__calc_digest()
422
423    def set_subkey(self, sign_algo, name, uuid, subkey_version, max_depth,
424                   name_size):
425        from cryptography.hazmat.primitives.asymmetric import rsa
426        import struct
427
428        self.subkey_name = name
429
430        subkey_key = load_asymmetric_key_img(self.inf)
431        if isinstance(subkey_key, rsa.RSAPrivateKey):
432            subkey_pkey = subkey_key.public_key()
433        else:
434            subkey_pkey = subkey_key
435
436        if max_depth is None:
437            if hasattr(self, 'previous_max_depth'):
438                if self.previous_max_depth <= 0:
439                    logger.error('Max depth of previous subkey is {}, '
440                                 .format(self.previous_max_depth) +
441                                 'cannot use a smaller value')
442                    sys.exit(1)
443
444                max_depth = self.previous_max_depth - 1
445            else:
446                max_depth = 0
447        else:
448            if (hasattr(self, 'previous_max_depth') and
449                    max_depth >= getattr(self, 'previous_max_depth')):
450                logger.error('Max depth of previous subkey is {} '
451                             .format(self.previous_max_depth) +
452                             'and the next value must be smaller')
453                sys.exit(1)
454
455        def int_to_bytes(x: int) -> bytes:
456            return x.to_bytes((x.bit_length() + 8) // 8, 'big')
457
458        n_bytes = int_to_bytes(subkey_pkey.public_numbers().n)
459        e_bytes = int_to_bytes(subkey_pkey.public_numbers().e)
460        attrs_end_offs = 16 + 5 * 4 + 2 * 3 * 4
461        shdr_subkey = struct.pack('<IIIIIIIIIII',
462                                  name_size, subkey_version,
463                                  max_depth, sig_tee_alg[sign_algo], 2,
464                                  TEE_ATTR_RSA_MODULUS,
465                                  attrs_end_offs, len(n_bytes),
466                                  TEE_ATTR_RSA_PUBLIC_EXPONENT,
467                                  attrs_end_offs + len(n_bytes),
468                                  len(e_bytes))
469        self.img = uuid.bytes + shdr_subkey + n_bytes + e_bytes
470        self.__pack_img(SHDR_SUBKEY, sign_algo)
471        self.img_digest = self.__calc_digest()
472
473    def parse(self):
474        from cryptography.hazmat.primitives.asymmetric import rsa
475        import struct
476
477        offs = 0
478        self.shdr = self.inf[offs:offs + SHDR_SIZE]
479        [magic, img_type, img_size, algo_value, hash_size,
480         sig_size] = struct.unpack('<IIIIHH', self.shdr)
481        offs += SHDR_SIZE
482
483        if magic != SHDR_MAGIC:
484            raise Exception("Unexpected magic: 0x{:08x}".format(magic))
485
486        if algo_value not in sig_tee_alg.values():
487            raise Exception('Unrecognized algorithm: 0x{:08x}'
488                            .format(algo_value))
489        self.sig_algo = value_to_key(sig_tee_alg, algo_value)
490
491        if hash_size != self.hash_size:
492            raise Exception("Unexpected digest len: {}".format(hash_size))
493
494        self.img_digest = self.inf[offs:offs + hash_size]
495        offs += hash_size
496        self.sig = self.inf[offs:offs + sig_size]
497        offs += sig_size
498
499        if img_type == SHDR_BOOTSTRAP_TA or img_type == SHDR_ENCRYPTED_TA:
500            self.ta_uuid = self.inf[offs:offs + UUID_SIZE]
501            offs += UUID_SIZE
502            self.ta_version = self.inf[offs:offs + 4]
503            offs += 4
504            if img_type == SHDR_ENCRYPTED_TA:
505                self.ehdr = self.inf[offs: offs + EHDR_SIZE]
506                offs += EHDR_SIZE
507                [enc_algo, flags, nonce_len,
508                 tag_len] = struct.unpack('<IIHH', self.ehdr)
509                if enc_algo not in enc_tee_alg.values():
510                    raise Exception('Unrecognized encrypt algorithm: 0x{:08x}'
511                                    .format(enc_algo))
512                if nonce_len != 12:
513                    raise Exception("Unexpected nonce len: {}"
514                                    .format(nonce_len))
515                self.nonce = self.inf[offs:offs + nonce_len]
516                offs += nonce_len
517
518                if tag_len != 16:
519                    raise Exception("Unexpected tag len: {}".format(tag_len))
520                self.tag = self.inf[offs:offs + tag_len]
521                offs += tag_len
522
523                self.ciphertext = self.inf[offs:]
524                if len(self.ciphertext) != img_size:
525                    raise Exception("Unexpected ciphertext size: ",
526                                    "got {}, expected {}"
527                                    .format(len(self.ciphertext), img_size))
528                self.img = self.ciphertext
529            else:
530                self.img = self.inf[offs:]
531                if len(self.img) != img_size:
532                    raise Exception("Unexpected img size: got {}, expected {}"
533                                    .format(len(self.img), img_size))
534        elif img_type == SHDR_SUBKEY:
535            subkey_offs = offs
536            self.uuid = self.inf[offs:offs + UUID_SIZE]
537            offs += UUID_SIZE
538            self.subkey_hdr = self.inf[offs:offs + SK_HDR_SIZE]
539            [self.name_size, self.subkey_version, self.max_depth, self.algo,
540             self.attr_count] = struct.unpack('<IIIII', self.subkey_hdr)
541            offs += len(self.subkey_hdr)
542            self.attr = self.inf[offs:offs + img_size -
543                                 UUID_SIZE - len(self.subkey_hdr)]
544            offs += len(self.attr)
545            self.name_img = self.inf[offs:offs + self.name_size]
546            offs += self.name_size
547            self.next_inf = self.inf[offs:]
548
549            def find_attr(attr):
550                if self.attr_count <= 0:
551                    return None
552                for n in range(self.attr_count):
553                    o = subkey_offs + UUID_SIZE + SK_HDR_SIZE + n * 12
554                    [attr_value, attr_offs,
555                     attr_len] = struct.unpack('<III', self.inf[o: o + 12])
556                    if attr_value == attr:
557                        o = subkey_offs + attr_offs
558                        return self.inf[o:o + attr_len]
559                return None
560
561            n_bytes = find_attr(TEE_ATTR_RSA_MODULUS)
562            e_bytes = find_attr(TEE_ATTR_RSA_PUBLIC_EXPONENT)
563            e = int.from_bytes(e_bytes, 'big')
564            n = int.from_bytes(n_bytes, 'big')
565            self.subkey_key = rsa.RSAPublicNumbers(e, n).public_key()
566
567            self.img = self.inf[subkey_offs:offs - self.name_size]
568            if len(self.img) != img_size:
569                raise Exception("Unexpected img size: got {}, expected {}"
570                                .format(len(self.img), img_size))
571        else:
572            raise Exception("Unsupported image type: {}".format(img_type))
573
574    def display(self):
575        import binascii
576        import struct
577        import uuid
578
579        def display_ta():
580            nonlocal offs
581            ta_uuid = self.inf[offs:offs + UUID_SIZE]
582            print(' struct shdr_bootstrap_ta')
583            print('  uuid:       {}'.format(uuid.UUID(bytes=ta_uuid)))
584            offs += UUID_SIZE
585            [ta_version] = struct.unpack('<I', self.inf[offs:offs + 4])
586            print('  ta_version: {}'.format(ta_version))
587
588            offs += 4
589            if img_type == SHDR_ENCRYPTED_TA:
590                ehdr = self.inf[offs: offs + EHDR_SIZE]
591                offs += EHDR_SIZE
592                [enc_algo, flags, nonce_len,
593                 tag_len] = struct.unpack('<IIHH', ehdr)
594
595                print(' struct shdr_encrypted_ta')
596                enc_algo_name = 'Unkown'
597                if enc_algo in enc_tee_alg.values():
598                    enc_algo_name = value_to_key(enc_tee_alg, enc_algo)
599                print('  enc_algo:   0x{:08x} ({})'
600                      .format(enc_algo, enc_algo_name))
601
602                if enc_algo not in enc_tee_alg.values():
603                    raise Exception('Unrecognized encrypt algorithm: 0x{:08x}'
604                                    .format(enc_algo))
605
606                flags_name = 'Unkown'
607                if flags in enc_key_type.values():
608                    flags_name = value_to_key(enc_key_type, flags)
609                print('  flags:      0x{:x} ({})'.format(flags, flags_name))
610
611                print('  iv_size:    {} (bytes)'.format(nonce_len))
612                if nonce_len != NONCE_SIZE:
613                    raise Exception("Unexpected nonce len: {}"
614                                    .format(nonce_len))
615                nonce = self.inf[offs:offs + nonce_len]
616                print('  iv:         {}'
617                      .format(binascii.hexlify(nonce).decode('ascii')))
618                offs += nonce_len
619
620                print('  tag_size:   {} (bytes)'.format(tag_len))
621                if tag_len != TAG_SIZE:
622                    raise Exception("Unexpected tag len: {}".format(tag_len))
623                tag = self.inf[offs:offs+tag_len]
624                print('  tag:        {}'
625                      .format(binascii.hexlify(tag).decode('ascii')))
626                offs += tag_len
627
628                ciphertext = self.inf[offs:]
629                print(' TA offset:  {} (0x{:x}) bytes'.format(offs, offs))
630                print(' TA size:    {} (0x{:x}) bytes'
631                      .format(len(ciphertext), len(ciphertext)))
632                if len(ciphertext) != img_size:
633                    raise Exception("Unexpected ciphertext size: ",
634                                    "got {}, expected {}"
635                                    .format(len(ciphertext), img_size))
636            else:
637                img = self.inf[offs:]
638                print(' TA offset:  {} (0x{:x}) bytes'.format(offs, offs))
639                print(' TA size:    {} (0x{:x}) bytes'
640                      .format(len(img), len(img)))
641                if len(img) != img_size:
642                    raise Exception("Unexpected img size: got {}, expected {}"
643                                    .format(len(img), img_size))
644            offs += img_size
645
646        offs = 0
647        while offs < len(self.inf):
648            if offs > 0:
649                # name_size is the previous subkey header
650                name_img = self.inf[offs:offs + name_size]
651                print('  next name:  "{}"'.format(name_img_to_str(name_img)))
652                offs += name_size
653                print('Next header at offset: {} (0x{:x})'
654                      .format(offs, offs))
655
656            shdr = self.inf[offs:offs + SHDR_SIZE]
657            [magic, img_type, img_size, algo_value, hash_size,
658             sig_size] = struct.unpack('<IIIIHH', shdr)
659            offs += SHDR_SIZE
660
661            if magic != SHDR_MAGIC:
662                Exception("Unexpected magic: 0x{:08x}".format(magic))
663
664            img_type_name = 'Unknown'
665            if img_type == SHDR_BOOTSTRAP_TA:
666                print('Bootstrap TA')
667                img_type_name = 'SHDR_BOOTSTRAP_TA'
668            if img_type == SHDR_ENCRYPTED_TA:
669                print('Encrypted TA')
670                img_type_name = 'SHDR_ENCRYPTED_TA'
671            if img_type == SHDR_SUBKEY:
672                print('Subkey')
673                img_type_name = 'SHDR_SUBKEY'
674
675            algo_name = 'Unknown'
676            if algo_value in sig_tee_alg.values():
677                algo_name = value_to_key(sig_tee_alg, algo_value)
678
679            print(' struct shdr')
680            print('  magic:      0x{:08x}'.format(magic))
681            print('  img_type:   {} ({})'.format(img_type, img_type_name))
682            print('  img_size:   {} bytes'.format(img_size))
683            print('  algo:       0x{:08x} ({})'.format(algo_value, algo_name))
684            print('  hash_size:  {} bytes'.format(hash_size))
685            print('  sig_size:   {} bytes'.format(sig_size))
686
687            if algo_value not in sig_tee_alg.values():
688                raise Exception('Unrecognized algorithm: 0x{:08x}'
689                                .format(algo_value))
690
691            if hash_size != self.hash_size:
692                raise Exception("Unexpected digest len: {}".format(hash_size))
693
694            img_digest = self.inf[offs:offs + hash_size]
695            print('  hash:       {}'
696                  .format(binascii.hexlify(img_digest).decode('ascii')))
697            offs += hash_size
698            sig = self.inf[offs:offs + sig_size]
699            offs += sig_size
700
701            if img_type == SHDR_BOOTSTRAP_TA or img_type == SHDR_ENCRYPTED_TA:
702                display_ta()
703            elif img_type == SHDR_SUBKEY:
704                img_uuid = self.inf[offs:offs + UUID_SIZE]
705                img_subkey = self.inf[offs + UUID_SIZE:
706                                      offs + UUID_SIZE + SK_HDR_SIZE]
707                [name_size, subkey_version, max_depth, algo,
708                 attr_count] = struct.unpack('<IIIII', img_subkey)
709                if algo not in sig_tee_alg.values():
710                    raise Exception('Unrecognized algorithm: 0x{:08x}'
711                                    .format(algo))
712                algo_name = value_to_key(sig_tee_alg, algo)
713                print(' struct shdr_subkey')
714                print('  uuid:       {}'.format(uuid.UUID(bytes=img_uuid)))
715                print('  name_size:  {}'.format(name_size))
716                print('  subkey_version: {}'.format(subkey_version))
717                print('  max_depth:  {}'.format(max_depth))
718                print('  algo:       0x{:08x} ({})'.format(algo, algo_name))
719                print('  attr_count: {}'.format(attr_count))
720                offs += img_size
721            else:
722                raise Exception("Unsupported image type: {}".format(img_type))
723
724    def decrypt_ta(self, enc_key):
725        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
726
727        cipher = AESGCM(bytes.fromhex(enc_key))
728        self.img = cipher.decrypt(self.nonce, self.ciphertext + self.tag, None)
729
730    def __get_padding(self):
731        from cryptography.hazmat.primitives.asymmetric import padding
732
733        if self.sig_algo == 'TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256':
734            pad = padding.PSS(mgf=padding.MGF1(self.chosen_hash),
735                              salt_length=self.hash_size)
736        elif self.sig_algo == 'TEE_ALG_RSASSA_PKCS1_V1_5_SHA256':
737            pad = padding.PKCS1v15()
738
739        return pad
740
741    def sign(self):
742        from cryptography.hazmat.primitives.asymmetric import utils
743        from cryptography.hazmat.primitives.asymmetric import rsa
744
745        if not isinstance(self.key, rsa.RSAPrivateKey):
746            logger.error('Provided key cannot be used for signing, ' +
747                         'please use offline-signing mode.')
748            sys.exit(1)
749        else:
750            self.sig = self.key.sign(self.img_digest, self.__get_padding(),
751                                     utils.Prehashed(self.chosen_hash))
752
753            if len(self.sig) != self.sig_size:
754                raise Exception(("Actual signature length is not equal to ",
755                                 "the computed one: {} != {}").
756                                format(len(self.sig), self.sig_size))
757
758    def add_signature(self, sigf):
759        import base64
760
761        with open(sigf, 'r') as f:
762            self.sig = base64.b64decode(f.read())
763
764        if len(self.sig) != self.sig_size:
765            raise Exception(("Actual signature length is not equal to ",
766                             "the expected one: {} != {}").
767                            format(len(self.sig), self.sig_size))
768
769    def verify_signature(self):
770        from cryptography.hazmat.primitives.asymmetric import utils
771        from cryptography.hazmat.primitives.asymmetric import rsa
772        from cryptography import exceptions
773
774        if isinstance(self.key, rsa.RSAPrivateKey):
775            pkey = self.key.public_key()
776        else:
777            pkey = self.key
778
779        try:
780            pkey.verify(self.sig, self.img_digest, self.__get_padding(),
781                        utils.Prehashed(self.chosen_hash))
782        except exceptions.InvalidSignature:
783            logger.error('Verification failed, ignoring given signature.')
784            sys.exit(1)
785
786    def verify_digest(self):
787        if self.img_digest != self.__calc_digest():
788            raise Exception('Hash digest does not match')
789
790    def verify_uuid(self, uuid):
791        if self.ta_uuid != uuid.bytes:
792            raise Exception('UUID does not match')
793
794    def add_subkey(self, subkey_file, name):
795        sk_image = BinaryImage(subkey_file, None)
796        self.subkey_img = sk_image.inf
797        sk_image.parse()
798        if not hasattr(sk_image, 'next_inf'):
799            logger.error('Invalid subkey file')
800            sys.exit(1)
801        while len(sk_image.next_inf) > 0:
802            sk_image = BinaryImage(sk_image.next_inf, None)
803            sk_image.parse()
804
805        if name is None:
806            name = ''
807        self.previous_max_depth = sk_image.max_depth
808        self.name_img = str.encode(name).ljust(sk_image.name_size, b'\0')
809
810    def write(self, outf):
811        with open(outf, 'wb') as f:
812            if hasattr(self, 'subkey_img'):
813                f.write(self.subkey_img)
814                f.write(self.name_img)
815            f.write(self.shdr)
816            f.write(self.img_digest)
817            f.write(self.sig)
818            if hasattr(self, 'ta_uuid'):
819                f.write(self.ta_uuid)
820                f.write(self.ta_version)
821            if hasattr(self, 'ehdr'):
822                f.write(self.ehdr)
823                f.write(self.nonce)
824                f.write(self.tag)
825                f.write(self.ciphertext)
826            else:
827                f.write(self.img)
828
829
830def load_ta_image(args):
831    ta_image = BinaryImage(args.inf, args.key)
832
833    if args.enc_key:
834        ta_image.encrypt_ta(args.enc_key, args.enc_key_type,
835                            args.algo, args.uuid, args.ta_version)
836    else:
837        ta_image.set_bootstrap_ta(args.algo, args.uuid, args.ta_version)
838
839    return ta_image
840
841
842def command_sign_enc(args):
843    ta_image = load_ta_image(args)
844    if args.subkey:
845        ta_image.add_subkey(args.subkey, args.name)
846    ta_image.sign()
847    ta_image.write(args.outf)
848    logger.info('Successfully signed application.')
849
850
851def command_sign_subkey(args):
852    image = BinaryImage(args.inf, args.key)
853    if args.subkey:
854        image.add_subkey(args.subkey, args.name)
855    image.set_subkey(args.algo, args.name, args.uuid, args.subkey_version,
856                     args.max_depth, args.name_size)
857    image.sign()
858    image.write(args.outf)
859    logger.info('Successfully signed subkey.')
860
861
862def command_digest(args):
863    import base64
864
865    ta_image = load_ta_image(args)
866    with open(args.digf, 'wb+') as digfile:
867        digfile.write(base64.b64encode(ta_image.img_digest))
868
869
870def command_stitch(args):
871    ta_image = load_ta_image(args)
872    ta_image.add_signature(args.sigf)
873    ta_image.verify_signature()
874    ta_image.write(args.outf)
875    logger.info('Successfully applied signature.')
876
877
878def command_verify(args):
879    import uuid
880
881    image = BinaryImage(args.inf, args.key)
882    next_uuid = None
883    max_depth = -1
884    while True:
885        image.parse()
886        if hasattr(image, 'subkey_hdr'):  # Subkey
887            print('Subkey UUID: {}'.format(uuid.UUID(bytes=image.uuid)))
888            image.verify_signature()
889            image.verify_digest()
890            if next_uuid:
891                if uuid.UUID(bytes=image.uuid) != next_uuid:
892                    raise Exception('UUID {} does not match {}'
893                                    .format(uuid.UUID(bytes=image.uuid),
894                                            next_uuid))
895            if max_depth >= 0:
896                if image.max_depth < 0 or image.max_depth >= max_depth:
897                    raise Exception('Invalid max_depth {} not less than {}'
898                                    .format(image.max_depth, max_depth))
899            max_depth = image.max_depth
900            if len(image.next_inf) == 0:
901                logger.info('Subkey is correctly verified.')
902                return
903            if image.name_size > 0:
904                next_uuid = uuid_v5_sha512(image.uuid,
905                                           name_img_to_str(image.name_img))
906            else:
907                next_uuid = image.uuid
908            image = BinaryImage(image.next_inf, image.subkey_key)
909        else:  # TA
910            print('TA UUID: {}'.format(uuid.UUID(bytes=image.ta_uuid)))
911            if next_uuid:
912                if uuid.UUID(bytes=image.ta_uuid) != next_uuid:
913                    raise Exception('UUID {} does not match {}'
914                                    .format(uuid.UUID(bytes=image.ta_uuid),
915                                            next_uuid))
916            if hasattr(image, 'ciphertext'):
917                if args.enc_key is None:
918                    logger.error('--enc-key needed to decrypt TA')
919                    sys.exit(1)
920                image.decrypt_ta(args.enc_key)
921            image.verify_signature()
922            image.verify_digest()
923            image.verify_uuid(args.uuid)
924            logger.info('Trusted application is correctly verified.')
925            return
926
927
928def command_display(args):
929    ta_image = BinaryImage(args.inf, None)
930    ta_image.display()
931
932
933def command_subkey_uuid(args):
934    import uuid
935
936    sk_image = BinaryImage(args.inf, None)
937    sk_image.parse()
938    if not hasattr(sk_image, 'next_inf'):
939        logger.error('Invalid subkey file')
940        sys.exit(1)
941    print('Subkey UUID: {}'.format(uuid.UUID(bytes=sk_image.uuid)))
942    while len(sk_image.next_inf) > 0:
943        sk_image = BinaryImage(sk_image.next_inf, None)
944        sk_image.parse()
945        print('Subkey UUID: {}'.format(uuid.UUID(bytes=sk_image.uuid)))
946    if args.name:
947        if len(args.name) > sk_image.name_size:
948            logger.error('Length of name ({}) '.format(len(args.name)) +
949                         'is larger than max name size ({})'
950                         .format(sk_image.name_size))
951            sys.exit(1)
952        print('Next subkey UUID: {}'
953              .format(uuid_v5_sha512(sk_image.uuid, args.name)))
954    else:
955        print('Next subkey UUID unchanged: {}'
956              .format(uuid.UUID(bytes=sk_image.uuid)))
957
958
959def main():
960    import logging
961    import os
962
963    global logger
964    logging.basicConfig()
965    logger = logging.getLogger(os.path.basename(__file__))
966
967    args = get_args()
968    args.func(args)
969
970
971if __name__ == "__main__":
972    main()
973