1#!/usr/bin/python3
2# Copyright (c) 2022, Arm Limited. All rights reserved.
3#
4# SPDX-License-Identifier: BSD-3-Clause
5
6#
7# Copyright 2022 The Hafnium Authors.
8#
9# Use of this source code is governed by a BSD-style
10# license that can be found in the LICENSE file or at
11# https://opensource.org/licenses/BSD-3-Clause.
12
13"""
14Script which generates a Secure Partition package.
15https://trustedfirmware-a.readthedocs.io/en/latest/components/secure-partition-manager.html#secure-partition-packages
16"""
17
18import argparse
19from collections import namedtuple
20import sys
21from shutil import copyfileobj
22import os
23
24HF_PAGE_SIZE = 0x1000 # bytes
25HEADER_ELEMENT_BYTES = 4 # bytes
26MANIFEST_IMAGE_SPLITTER=':'
27PM_OFFSET_DEFAULT = "0x1000"
28IMG_OFFSET_DEFAULT = "0x4000"
29
30def split_dtb_bin(i : str):
31    return i.split(MANIFEST_IMAGE_SPLITTER)
32
33def align_to_page(n):
34    return HF_PAGE_SIZE * \
35          (round(n / HF_PAGE_SIZE) + \
36           (1 if n % HF_PAGE_SIZE else 0))
37
38def to_bytes(value):
39    return int(value).to_bytes(HEADER_ELEMENT_BYTES, 'little')
40
41class SpPkg:
42    def __init__(self, pm_path : str, img_path : str, pm_offset: int,
43                 img_offset: int):
44        if not os.path.isfile(pm_path) or not os.path.isfile(img_path):
45            raise Exception(f"Parameters should be path.  \
46                              manifest: {pm_path}; img: {img_path}")
47        self.pm_path = pm_path
48        self.img_path = img_path
49        self._SpPkgHeader = namedtuple("SpPkgHeader",
50                             ("magic", "version",
51                              "pm_offset", "pm_size",
52                              "img_offset", "img_size"))
53
54        if pm_offset >= img_offset:
55            raise ValueError("pm_offset must be smaller than img_offset")
56
57        is_hfpage_aligned = lambda val : val % HF_PAGE_SIZE == 0
58        if not is_hfpage_aligned(pm_offset) or not is_hfpage_aligned(img_offset):
59           raise ValueError(f"Offsets provided need to be page aligned: pm-{pm_offset}, img-{img_offset}")
60
61        if img_offset - pm_offset < self.pm_size:
62            raise ValueError(f"pm_offset and img_offset do not fit the specified file:{pm_path})")
63
64        self.pm_offset = pm_offset
65        self.img_offset = img_offset
66
67    def __str__(self):
68        return \
69        f'''--SP package Info--
70        header:{self.header}
71        pm: {self.pm_path}
72        img: {self.img_path}
73        '''
74
75    @property
76    def magic(self):
77        return "SPKG".encode()
78
79    @property
80    def version(self):
81        return 0x2
82
83    @property
84    def pm_size(self):
85        return os.path.getsize(self.pm_path)
86
87    @property
88    def img_size(self):
89        return os.path.getsize(self.img_path)
90
91    @property
92    def header(self):
93        return self._SpPkgHeader(
94                self.magic,
95                self.version,
96                self.pm_offset,
97                self.pm_size,
98                self.img_offset,
99                self.img_size)
100
101    @property
102    def header_size(self):
103        return len(self._SpPkgHeader._fields)
104
105    def generate(self, f_out : str):
106        with open(f_out, "wb+") as output:
107            for h in self.header:
108                to_write = h if type(h) is bytes else to_bytes(h)
109                output.write(to_write)
110            output.seek(self.pm_offset)
111            with open(self.pm_path, "rb") as pm:
112                copyfileobj(pm, output)
113            output.seek(self.img_offset)
114            with open(self.img_path, "rb") as img:
115                copyfileobj(img, output)
116
117def Main():
118    parser = argparse.ArgumentParser()
119    parser.add_argument("-i", required=True,
120                        help="path to partition's image and manifest separated by a colon.")
121    parser.add_argument("--pm-offset", required=False, default=PM_OFFSET_DEFAULT,
122                        help="set partitition manifest offset.")
123    parser.add_argument("--img-offset", required=False, default=IMG_OFFSET_DEFAULT,
124                        help="set partition image offset.")
125    parser.add_argument("-o", required=True, help="set output file path.")
126    parser.add_argument("-v", required=False, action="store_true",
127                        help="print package information.")
128    args = parser.parse_args()
129
130    if not os.path.exists(os.path.dirname(args.o)):
131        raise Exception("Provide a valid output file path!\n")
132
133    image_path, manifest_path = split_dtb_bin(args.i)
134    pm_offset = int(args.pm_offset, 0)
135    img_offset = int(args.img_offset, 0)
136    pkg = SpPkg(manifest_path, image_path, pm_offset, img_offset)
137    pkg.generate(args.o)
138
139    if args.v is True:
140        print(pkg)
141
142    return 0
143
144if __name__ == "__main__":
145    sys.exit(Main())
146