1#-------------------------------------------------------------------------------
2# Copyright (c) 2022-2023, Arm Limited. All rights reserved.
3#
4# SPDX-License-Identifier: BSD-3-Clause
5#
6#-------------------------------------------------------------------------------
7
8import argparse
9import hashlib
10from functools import reduce
11from operator import add
12from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
13import struct
14import secrets
15
16sic_page_size = 1024
17sic_line_size = 16
18
19def struct_pack(objects, pad_to=0):
20    defstring = "<"
21    for obj in objects:
22        defstring += str(len(obj)) + "s"
23
24    size = struct.calcsize(defstring)
25    if size < pad_to:
26        defstring += str(pad_to - size) + "x"
27
28    return (bytes(struct.pack(defstring, *objects)))
29
30def chunk_bytes(x, n):
31    return [x[i:i+n] for i in range(0, len(x), n)]
32
33def round_up(x, boundary):
34    return ((x + (boundary - 1)) // boundary) * boundary
35
36parser = argparse.ArgumentParser()
37parser.add_argument("--input_image", help="the image to create table from", required=True)
38parser.add_argument("--encrypt_key_file", help="Key to encrypt image with", required=False)
39parser.add_argument("--image_version", help="Version of the image", required=True)
40parser.add_argument("--table_output_file", help="table output file", required=True)
41parser.add_argument("--encrypted_image_output_file", help="encrupted image output file", required=True)
42args = parser.parse_args()
43
44with open(args.input_image, "rb") as in_file:
45    image = in_file.read()
46
47if args.encrypt_key_file is not None:
48    with open(args.encrypt_key, "rb") as in_file:
49        encrypt_key = in_file.read()
50else:
51    encrypt_key = bytearray([0xfc, 0x57, 0x01, 0xdc, 0x61, 0x35, 0xe1, 0x32,
52                             0x38, 0x47, 0xbd, 0xc4, 0x0f, 0x04, 0xd2, 0xe5,
53                             0xbe, 0xe5, 0x83, 0x3b, 0x23, 0xc2, 0x9f, 0x93,
54                             0x59, 0x3d, 0x00, 0x01, 0x8c, 0xfa, 0x99, 0x94,])
55
56
57fw_version_bytes = int(args.image_version, 0).to_bytes(4, 'little')
58nonce_bytes = secrets.token_bytes(8)
59
60# We need to pad the image to the authentication page size
61image = struct_pack([image], round_up(len(image), sic_page_size))
62
63# The SIC uses a non-standard counter construction, so we need to do this
64# manually
65enc_image = []
66line_idx = 0
67for chunk in chunk_bytes(image, sic_line_size):
68    counter_val = struct_pack([line_idx.to_bytes(4, 'little'), fw_version_bytes, nonce_bytes[4:], nonce_bytes[:4]], pad_to=16)
69    line_idx += 1
70    cipher = Cipher(algorithms.AES(encrypt_key), modes.CTR(counter_val))
71    enc_image.append(cipher.encryptor().update(chunk))
72enc_image = reduce(add, enc_image, b"")
73
74htr = reduce(add, map(lambda x:hashlib.sha256(x).digest(), chunk_bytes(enc_image, sic_page_size)), b"")
75
76table = struct_pack([
77            fw_version_bytes,
78            nonce_bytes,
79            len(htr).to_bytes(4, 'little'),
80            htr
81            ])
82
83with open(args.table_output_file, "wb") as out_file:
84    out_file.write(table)
85
86with open(args.encrypted_image_output_file, "wb") as out_file:
87    out_file.write(enc_image)
88