1#-------------------------------------------------------------------------------
2# Copyright (c) 2021-2023, Arm Limited. All rights reserved.
3#
4# SPDX-License-Identifier: BSD-3-Clause
5#
6#-------------------------------------------------------------------------------
7
8import hashlib
9from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
10from cryptography.hazmat.primitives import cmac
11from cryptography.hazmat.backends import default_backend
12import secrets
13import argparse
14import os
15import sys
16sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), "../../../bl2/ext/mcuboot/scripts"))
17import macro_parser
18import struct
19import pyhsslms
20
21def struct_pack(objects, pad_to=0):
22    defstring = "<"
23    for obj in objects:
24        defstring += str(len(obj)) + "s"
25
26    size = struct.calcsize(defstring)
27    if size < pad_to:
28        defstring += str(pad_to - size) + "x"
29
30    return (bytes(struct.pack(defstring, *objects)))
31
32def parse_version(version_string):
33    version = [0, 0, 0, 0]
34    split = version_string.split("+")
35    if len(split) > 1:
36        version[3] = int(split[1])
37    split = split[0].split(".")
38    for i in range(len(split)):
39        version[i] = int(split[i])
40
41    return struct_pack([version[0].to_bytes(1, "little"),
42                        version[1].to_bytes(1, "little"),
43                        version[2].to_bytes(2, "little"),
44                        version[3].to_bytes(4, "little")])
45
46def derive_encryption_key(security_counter):
47    with open(args.encrypt_key_file, "rb") as encrypt_key_file:
48        encrypt_key = encrypt_key_file.read()
49
50    output_key = bytes(0);
51    # The KDF outputs 16 bytes per iteration, so we need 2 for an AES-256 key
52    for i in range(2):
53        state = struct_pack([(i + 1).to_bytes(4, byteorder='little'),
54                             # C keeps the null byte, python removes it, so we add
55                             # it back manually.
56                             "BL2_DECRYPTION_KEY".encode('ascii') + bytes(1),
57                             bytes(1), security_counter,
58                             (32).to_bytes(4, byteorder='little')])
59        c = cmac.CMAC(algorithms.AES(encrypt_key))
60        c.update(state)
61        output_key += c.finalize()
62    return output_key
63
64
65
66def sign_binary_blob(blob):
67    priv_key = pyhsslms.HssLmsPrivateKey(args.sign_key_file)
68    # Remove the first 4 bytes since it's HSS info
69    sig = priv_key.sign(blob)[4:]
70    if (len(sig) != 1452):
71        raise Exception
72    return sig
73
74def hash_binary_blob(blob):
75   hash = hashlib.sha256()
76   hash.update(blob)
77   return hash.digest()
78
79def encrypt_binary_blob(blob, counter_val, encrypt_key):
80    cipher = Cipher(algorithms.AES(encrypt_key), modes.CTR(counter_val))
81    return cipher.encryptor().update(blob)
82
83parser = argparse.ArgumentParser()
84parser.add_argument("--input_file", help="the image to process", required=True)
85parser.add_argument("--img_version", help="version of the image", required=True)
86parser.add_argument("--img_security_counter", help="Secuity counter value for the image", required=True)
87parser.add_argument("--encrypt_key_file", help="encryption key file", required=True)
88parser.add_argument("--sign_key_file", help="signing key file", required=False)
89parser.add_argument("--img_output_file", help="image output file", required=True)
90parser.add_argument("--hash_output_file", help="hash output file", required=False)
91parser.add_argument("--signing_layout_file", help="signing layout file", required=True)
92parser.add_argument("--header_size", help="size of the header", required=True)
93args = parser.parse_args()
94
95with open(args.input_file, "rb") as in_file:
96    bl2_code = in_file.read()
97
98counter_val = secrets.token_bytes(12) + int(0).to_bytes(4, 'little')
99
100version = parse_version(args.img_version)
101
102bl2_partition_size = macro_parser.evaluate_macro(args.signing_layout_file,
103                                    ".*(RE_BL2_BIN_SIZE) = *(.*)",
104                                    1, 2, True)['RE_BL2_BIN_SIZE']
105
106plaintext = struct_pack([
107    int("0xDEADBEEF", 16).to_bytes(4, 'little'),
108    int(0).to_bytes(int(args.header_size, 0) - (1452 + 16 + 8 + 4 + 4), 'little'),
109    bl2_code,
110    ],
111    pad_to=bl2_partition_size - (1452 + 16 + 8 + 4))
112
113encrypt_key = derive_encryption_key(int(args.img_security_counter, 16).to_bytes(4, 'little'))
114ciphertext = encrypt_binary_blob(plaintext, counter_val, encrypt_key)
115
116data_to_sign = struct_pack([
117    version,
118    int(args.img_security_counter, 0).to_bytes(4, 'little'),
119    plaintext,
120    ])
121
122hash = hash_binary_blob(data_to_sign)
123sig = sign_binary_blob(data_to_sign)
124
125image = struct_pack([
126    counter_val,
127    sig,
128    version,
129    int(args.img_security_counter, 0).to_bytes(4, 'little'),
130    ciphertext,
131    ])
132
133if len(image) > bl2_partition_size:
134    print("Error: Signed image size {} exceeds BL2 partition size {}"
135          .format(len(image), bl2_partition_size))
136    exit(1)
137
138with open(args.img_output_file, "wb") as img_out_file:
139    img_out_file.write(image)
140
141with open(args.hash_output_file, "wb") as hash_out_file:
142    hash_out_file.write(hash)
143