1#!/usr/bin/env python
2#
3# check_sizes.py is a tool run by the ESP-IDF build system
4# to check a particular binary fits in the available partitions of
5# a particular type/subtype. Can be used to check if the app binary fits in
6# all available app partitions, for example.
7#
8# (Can also check if the bootloader binary fits before the partition table.)
9#
10# Copyright 2020 Espressif Systems (Shanghai) PTE LTD
11#
12# Licensed under the Apache License, Version 2.0 (the "License");
13# you may not use this file except in compliance with the License.
14# You may obtain a copy of the License at
15#
16#     http://www.apache.org/licenses/LICENSE-2.0
17#
18# Unless required by applicable law or agreed to in writing, software
19# distributed under the License is distributed on an "AS IS" BASIS,
20# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
21# See the License for the specific language governing permissions and
22# limitations under the License.
23from __future__ import division, print_function, unicode_literals
24
25import argparse
26import io  # noqa: F401 # pylint: disable=unused-import
27import os
28import sys
29
30try:
31    from typing import IO  # noqa: F401 # pylint: disable=unused-import
32except ImportError:
33    pass  # used for type hinting only
34
35import gen_esp32part
36from gen_esp32part import PartitionTable, get_ptype_as_int, get_subtype_as_int
37
38allow_failures = False
39
40
41def _file_size(f):   # type: (IO) -> int
42    before = f.tell()
43    f.seek(0, 2)  # seek to end
44    result = f.tell()
45    f.seek(before)
46    return result
47
48
49def _fail(msg):   # type: (str) -> None
50    if allow_failures:
51        print('Warning: {}'.format(msg))
52    else:
53        raise SystemExit('Error: {}'.format(msg))
54
55
56def check_bootloader(partition_table_offset, bootloader_offset, binary_file):  # type: (int, int, IO) -> None
57    max_size = partition_table_offset - bootloader_offset
58    bootloader_size = _file_size(binary_file)
59    if bootloader_size > max_size:
60        msg = ('Bootloader binary size {:#x} bytes is too large for partition table offset {:#02x}. ' +
61               'Bootloader binary can be maximum {:#x} ({}) bytes unless the partition table offset ' +
62               'is increased in the Partition Table section of the project configuration menu.').format(
63                   bootloader_size, partition_table_offset, max_size, max_size)
64        _fail(msg)
65    free_size = max_size - bootloader_size
66    print('Bootloader binary size {:#x} bytes. {:#x} bytes ({}%) free.'.format(
67          bootloader_size, free_size, round(free_size * 100 / max_size)))
68
69
70def check_partition(ptype, subtype, partition_table_file, bin_file):  # type: (str, str, io.IOBase, IO) -> None
71    table, _ = PartitionTable.from_file(partition_table_file)
72    ptype_str = str(ptype)
73    ptype = get_ptype_as_int(ptype)
74    partitions = [p for p in table if p.type == ptype]
75
76    if subtype is not None:
77        ptype_str += ' ({})'.format(subtype)
78        subtype = get_subtype_as_int(ptype, subtype)
79        partitions = [p for p in partitions if p.subtype == subtype]
80
81    if len(partitions) == 0:
82        print('WARNING: Partition table does not contain any partitions matching {}'.format(ptype_str))
83        return
84
85    bin_name = os.path.basename(bin_file.name)
86    bin_size = _file_size(bin_file)
87    smallest_size = min(p.size for p in partitions)
88    if smallest_size >= bin_size:
89        free_size = smallest_size - bin_size
90        print('{} binary size {:#x} bytes. Smallest {} partition is {:#x} bytes. {:#x} bytes ({}%) free.'.format(
91            bin_name, bin_size, ptype_str, smallest_size, free_size, round(free_size * 100 / smallest_size)))
92        return
93
94    too_small_partitions = [p for p in partitions if p.size < bin_size]
95    if len(partitions) == 1:
96        msg = '{} partition is'.format(ptype_str)
97    elif len(partitions) == len(too_small_partitions):
98        msg = 'All {} partitions are'.format(ptype_str)
99    else:
100        msg = '{}/{} {} partitions are'.format(len(too_small_partitions), len(partitions), ptype_str)
101    msg += ' too small for binary {} size {:#x}:'.format(bin_name, bin_size)
102    for p in too_small_partitions:
103        msg += '\n  - {} (overflow {:#x})'.format(p, bin_size - p.size)
104    if not allow_failures and len(partitions) == len(too_small_partitions):
105        # if some partitions can fit the binary then just print a warning
106        raise SystemExit('Error: ' + msg)
107    else:
108        print('Warning: ' + msg)
109
110
111def main():  # type: () -> None
112    global allow_failures  # pylint: disable=global-statement
113
114    parser = argparse.ArgumentParser(description='Check binary sizes against partition table entries')
115    parser.add_argument('--target', choices=['esp32', 'esp32s2'])
116    parser.add_argument('--allow_failures', action='store_true', help='If true, failures will print warnings but not exit with an error')
117    parser.add_argument('--offset', '-o', help='Set partition table offset', default='0x8000')
118
119    subparsers = parser.add_subparsers(dest='check_target',
120                                       help='Type of binary to check against partition table layout')
121    sp_bootloader = subparsers.add_parser('bootloader')
122    sp_bootloader.add_argument('bootloader_offset', help='Hex offset of bootloader in flash')
123    sp_bootloader.add_argument('bootloader_binary', type=argparse.FileType('rb'), help='Bootloader binary (.bin) file from build output')
124
125    sp_part = subparsers.add_parser('partition')
126    sp_part.add_argument('--type', type=str, help='Check the file size against all partitions of this type.', required=True)
127    sp_part.add_argument('--subtype', type=str, help='Optional, only check the file size against all partitions of this subtype.')
128    sp_part.add_argument('partition_table', type=argparse.FileType('rb'), help='Partition table file')
129    sp_part.add_argument('binary', type=argparse.FileType('rb'), help='Binary file which will have the size checked')
130
131    args = parser.parse_args()
132
133    gen_esp32part.quiet = True
134
135    args.offset = int(args.offset, 0)
136    gen_esp32part.offset_part_table = args.offset
137
138    if args.check_target is None:  # add_subparsers only has a 'required' argument since Python 3
139        parser.print_help()
140        sys.exit(1)
141    if args.check_target == 'bootloader':
142        check_bootloader(args.offset, int(args.bootloader_offset, 0), args.bootloader_binary)
143    else:
144        check_partition(args.type, args.subtype, args.partition_table, args.binary)
145
146
147if __name__ == '__main__':
148    main()
149