1#-------------------------------------------------------------------------------
2# Copyright (c) 2023, Arm Limited. All rights reserved.
3#
4# SPDX-License-Identifier: BSD-3-Clause
5#
6#-------------------------------------------------------------------------------
7
8import argparse
9import struct
10from elftools.elf.elffile import ELFFile
11
12
13def struct_pack(objects, pad_to=0):
14    defstring = "<"
15    for obj in objects:
16        defstring += str(len(obj)) + "s"
17
18    size = struct.calcsize(defstring)
19    if size < pad_to:
20        defstring += str(pad_to - size) + "x"
21    elif size > pad_to and pad_to != 0:
22        print("Error padding struct of size {} to {}".format(size, pad_to))
23        exit(1)
24
25    return (bytes(struct.pack(defstring, *objects)))
26
27
28parser = argparse.ArgumentParser()
29parser.add_argument("--provisioning_bundle_axf",
30                    help="the input provisioning bundle elf/axf",
31                    required=True)
32parser.add_argument("--magic",
33                    help="the magic constant to insert at the start and end",
34                    required=True)
35parser.add_argument("--bl1_2_padded_hash_input_file",
36                    help="the hash of the final bl1_2 image",
37                    required=False)
38parser.add_argument("--bl1_2_input_file",
39                    help="the final bl1_2 image",
40                    required=False)
41parser.add_argument("--code_pad_size",
42                    help="size to pad the code section",
43                    required=True)
44parser.add_argument("--values_pad_size",
45                    help="size to pad the values section",
46                    required=True)
47parser.add_argument("--data_pad_size",
48                    help="size to pad the data section",
49                    required=True)
50parser.add_argument("--bundle_output_file",
51                    help="bundle output file",
52                    required=False)
53args = parser.parse_args()
54
55elffile = ELFFile(open(args.provisioning_bundle_axf, 'rb'))
56
57values_section = elffile.get_section_by_name("VALUES")
58rodata_section = elffile.get_section_by_name("RO_DATA")
59rwdata_section = elffile.get_section_by_name("RW_DATA")
60code_section = elffile.get_section_by_name("CODE")
61
62if code_section is not None:
63    code = code_section.data()
64else:
65    print("provisioning_bundle's code sections is mandatory")
66    exit(1)
67
68if rwdata_section is not None:
69    rwdata = rwdata_section.data()
70else:
71    rwdata = bytes(0)
72
73if rodata_section is not None:
74    rodata = rodata_section.data()
75else:
76    rodata = bytes(0)
77
78if values_section is not None:
79    values = values_section.data()
80else:
81    print("provisioning_bundle's values sections is mandatory")
82    exit(1)
83
84
85if args.bl1_2_padded_hash_input_file:
86    with open(args.bl1_2_padded_hash_input_file, "rb") as in_file:
87        bl1_2_padded_hash = in_file.read()
88else:
89    bl1_2_padded_hash = bytes(0)
90
91if args.bl1_2_input_file:
92    with open(args.bl1_2_input_file, "rb") as in_file:
93        bl1_2 = in_file.read()
94else:
95    bl1_2 = bytes(0)
96
97patch_bundle = struct_pack([
98    bl1_2_padded_hash,
99    bl1_2,
100])
101
102code = struct_pack([code],
103                   pad_to=int(args.code_pad_size, 0))
104values = struct_pack([patch_bundle, values[len(patch_bundle):]],
105                     pad_to=int(args.values_pad_size, 0))
106data = struct_pack([values, rwdata, rodata],
107                   pad_to=int(args.values_pad_size, 0)+int(args.data_pad_size, 0))
108
109bundle = struct_pack([
110    int(args.magic, 16).to_bytes(4, 'little'),
111    code,
112    data,
113    bytes(16), # Replace with GCM TAG
114    int(args.magic, 16).to_bytes(4, 'little'),
115])
116
117with open(args.bundle_output_file, "wb") as out_file:
118    out_file.write(bundle)
119