1#! /usr/bin/env python3
2#
3# Copyright 2017 Linaro Limited
4# Copyright (c) 2017-2019, Arm Limited.
5#
6# Licensed under the Apache License, Version 2.0 (the "License");
7# you may not use this file except in compliance with the License.
8# You may obtain a copy of the License at
9#
10#     http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing, software
13# distributed under the License is distributed on an "AS IS" BASIS,
14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15# See the License for the specific language governing permissions and
16# limitations under the License.
17
18"""
19Assemble multiple images into a single image that can be flashed on the device.
20"""
21
22import argparse
23import errno
24import io
25import re
26import os
27import shutil
28import macro_parser
29
30offset_re = re.compile(r"^\s*RE_([0-9A-Z_]+)_IMAGE_OFFSET\s*=\s*(.*)")
31size_re = re.compile(r"^\s*RE_([0-9A-Z_]+)_IMAGE_MAX_SIZE\s*=\s*(.*)")
32
33class Assembly():
34    def __init__(self, layout_path, output):
35        self.output = output
36        self.layout_path = layout_path
37        self.find_slots()
38        try:
39            os.unlink(output)
40        except OSError as e:
41            if e.errno != errno.ENOENT:
42                raise
43
44    def find_slots(self):
45        offsets = {}
46        sizes = {}
47
48        offsets = macro_parser.evaluate_macro(self.layout_path, offset_re, 1, 2)
49        sizes = macro_parser.evaluate_macro(self.layout_path, size_re, 1, 2)
50
51        if 'SECURE' not in offsets:
52            raise Exception("Image config does not have secure partition")
53
54        if 'NON_SECURE' not in offsets:
55            raise Exception("Image config does not have non-secure partition")
56
57        self.offsets = offsets
58        self.sizes = sizes
59
60    def add_image(self, source, partition):
61        with open(self.output, 'ab') as ofd:
62            ofd.seek(0, os.SEEK_END)
63            pos = ofd.tell()
64            if pos > self.offsets[partition]:
65                raise Exception("Partitions not in order, unsupported")
66            if pos < self.offsets[partition]:
67                ofd.write(b'\xFF' * (self.offsets[partition] - pos))
68            statinfo = os.stat(source)
69            if statinfo.st_size > self.sizes[partition]:
70                raise Exception("Image {} is too large for partition".format(source))
71            with open(source, 'rb') as rfd:
72                shutil.copyfileobj(rfd, ofd, 0x10000)
73
74def main():
75    parser = argparse.ArgumentParser()
76
77    parser.add_argument('-l', '--layout', required=True,
78            help='Location of the file that contains preprocessed macros')
79    parser.add_argument('-s', '--secure', required=True,
80            help='Unsigned secure image')
81    parser.add_argument('-n', '--non_secure',
82            help='Unsigned non-secure image')
83    parser.add_argument('-o', '--output', required=True,
84            help='Filename to write full image to')
85
86    args = parser.parse_args()
87    output = Assembly(args.layout, args.output)
88
89    output.add_image(args.secure, "SECURE")
90    output.add_image(args.non_secure, "NON_SECURE")
91
92if __name__ == '__main__':
93    main()
94