1#-------------------------------------------------------------------------------
2# Copyright (c) 2021-2024, Arm Limited. All rights reserved.
3#
4# SPDX-License-Identifier: BSD-3-Clause
5#
6#-------------------------------------------------------------------------------
7
8import struct
9import secrets
10from cryptography.hazmat.primitives.ciphers.aead import AESCCM
11from cryptography.hazmat.primitives import cmac
12from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
13import hashlib
14
15def patch_binary(binary, patch_bundle, offset):
16    return binary[:offset] + patch_bundle + binary[offset + len(patch_bundle):]
17
18def struct_pack(objects, pad_to=0):
19    defstring = "<"
20    for obj in objects:
21        defstring += str(len(obj)) + "s"
22
23    size = struct.calcsize(defstring)
24    if size < pad_to:
25        defstring += str(pad_to - size) + "x"
26    elif size > pad_to and pad_to != 0:
27        print("Error padding struct of size {} to {}".format(size, pad_to))
28        exit(1);
29
30    return (bytes(struct.pack(defstring, *objects)))
31
32def encrypt_bundle(code, code_pad, values, values_pad, data,
33                   data_and_values_pad, magic, input_key):
34    code = struct_pack([code], pad_to=code_pad)
35    values = struct_pack([values], pad_to=values_pad)
36    data = struct_pack([values, data], pad_to=data_and_values_pad)
37
38    ccm_iv = secrets.token_bytes(12)
39
40    to_auth = struct_pack([
41        magic.to_bytes(8, 'little'),
42    ])
43
44    to_encrypt = struct_pack([
45        code,
46        data,
47    ])
48
49    encrypted_data = AESCCM(input_key, tag_length=16).encrypt(ccm_iv, to_encrypt, to_auth)
50    tag = encrypted_data[-16:]
51    encrypted_data = encrypted_data[:-16]
52
53    bundle = struct_pack([
54        to_auth,
55        encrypted_data,
56        ccm_iv,
57        tag,
58        magic.to_bytes(8, 'little'),
59    ])
60
61    return bundle
62