1#!/usr/bin/env python
2#
3# parttool is used to perform partition level operations - reading,
4# writing, erasing and getting info about the partition.
5#
6# Copyright 2018 Espressif Systems (Shanghai) PTE LTD
7#
8# Licensed under the Apache License, Version 2.0 (the "License");
9# you may not use this file except in compliance with the License.
10# You may obtain a copy of the License at
11#
12#     http:#www.apache.org/licenses/LICENSE-2.0
13#
14# Unless required by applicable law or agreed to in writing, software
15# distributed under the License is distributed on an "AS IS" BASIS,
16# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17# See the License for the specific language governing permissions and
18# limitations under the License.
19from __future__ import division, print_function
20
21import argparse
22import os
23import re
24import subprocess
25import sys
26import tempfile
27
28import gen_esp32part as gen
29
30__version__ = '2.0'
31
32COMPONENTS_PATH = os.path.expandvars(os.path.join('$IDF_PATH', 'components'))
33ESPTOOL_PY = os.path.join(COMPONENTS_PATH, 'esptool_py', 'esptool', 'esptool.py')
34
35PARTITION_TABLE_OFFSET = 0x8000
36
37
38quiet = False
39
40
41def status(msg):
42    if not quiet:
43        print(msg)
44
45
46class _PartitionId():
47
48    def __init__(self, name=None, p_type=None, subtype=None, part_list=None):
49        self.name = name
50        self.type = p_type
51        self.subtype = subtype
52        self.part_list = part_list
53
54
55class PartitionName(_PartitionId):
56
57    def __init__(self, name):
58        _PartitionId.__init__(self, name=name)
59
60
61class PartitionType(_PartitionId):
62
63    def __init__(self, p_type, subtype, part_list=None):
64        _PartitionId.__init__(self, p_type=p_type, subtype=subtype, part_list=part_list)
65
66
67PARTITION_BOOT_DEFAULT = _PartitionId()
68
69
70class ParttoolTarget():
71
72    def __init__(self, port=None, baud=None, partition_table_offset=PARTITION_TABLE_OFFSET, partition_table_file=None,
73                 esptool_args=[], esptool_write_args=[], esptool_read_args=[], esptool_erase_args=[]):
74        self.port = port
75        self.baud = baud
76
77        gen.offset_part_table = partition_table_offset
78
79        def parse_esptool_args(esptool_args):
80            results = list()
81            for arg in esptool_args:
82                pattern = re.compile(r'(.+)=(.+)')
83                result = pattern.match(arg)
84                try:
85                    key = result.group(1)
86                    value = result.group(2)
87                    results.extend(['--' + key, value])
88                except AttributeError:
89                    results.extend(['--' + arg])
90            return results
91
92        self.esptool_args = parse_esptool_args(esptool_args)
93        self.esptool_write_args = parse_esptool_args(esptool_write_args)
94        self.esptool_read_args = parse_esptool_args(esptool_read_args)
95        self.esptool_erase_args = parse_esptool_args(esptool_erase_args)
96
97        if partition_table_file:
98            partition_table = None
99            with open(partition_table_file, 'rb') as f:
100                input_is_binary = (f.read(2) == gen.PartitionDefinition.MAGIC_BYTES)
101                f.seek(0)
102                if input_is_binary:
103                    partition_table = gen.PartitionTable.from_binary(f.read())
104
105            if partition_table is None:
106                with open(partition_table_file, 'r') as f:
107                    f.seek(0)
108                    partition_table = gen.PartitionTable.from_csv(f.read())
109        else:
110            temp_file = tempfile.NamedTemporaryFile(delete=False)
111            temp_file.close()
112
113            try:
114                self._call_esptool(['read_flash', str(partition_table_offset), str(gen.MAX_PARTITION_LENGTH), temp_file.name])
115                with open(temp_file.name, 'rb') as f:
116                    partition_table = gen.PartitionTable.from_binary(f.read())
117            finally:
118                os.unlink(temp_file.name)
119
120        self.partition_table = partition_table
121
122    # set `out` to None to redirect the output to the STDOUT
123    # otherwise set `out` to file descriptor
124    # beware that the method does not close the file descriptor
125    def _call_esptool(self, args, out=None):
126        esptool_args = [sys.executable, ESPTOOL_PY] + self.esptool_args
127
128        if self.port:
129            esptool_args += ['--port', self.port]
130
131        if self.baud:
132            esptool_args += ['--baud', str(self.baud)]
133
134        esptool_args += args
135
136        print('Running %s...' % (' '.join(esptool_args)))
137        try:
138            subprocess.check_call(esptool_args, stdout=out, stderr=subprocess.STDOUT)
139        except subprocess.CalledProcessError as e:
140            print('An exception: **', str(e), '** occurred in _call_esptool.', file=out)
141            raise e
142
143    def get_partition_info(self, partition_id):
144        partition = None
145
146        if partition_id.name:
147            partition = self.partition_table.find_by_name(partition_id.name)
148        elif partition_id.type and partition_id.subtype:
149            partition = list(self.partition_table.find_by_type(partition_id.type, partition_id.subtype))
150            if not partition_id.part_list:
151                partition = partition[0]
152        else:  # default boot partition
153            search = ['factory'] + ['ota_{}'.format(d) for d in range(16)]
154            for subtype in search:
155                partition = next(self.partition_table.find_by_type('app', subtype), None)
156                if partition:
157                    break
158
159        if not partition:
160            raise Exception('Partition does not exist')
161
162        return partition
163
164    def erase_partition(self, partition_id):
165        partition = self.get_partition_info(partition_id)
166        self._call_esptool(['erase_region', str(partition.offset),  str(partition.size)] + self.esptool_erase_args)
167
168    def read_partition(self, partition_id, output):
169        partition = self.get_partition_info(partition_id)
170        self._call_esptool(['read_flash', str(partition.offset), str(partition.size), output] + self.esptool_read_args)
171
172    def write_partition(self, partition_id, input):
173        self.erase_partition(partition_id)
174
175        partition = self.get_partition_info(partition_id)
176
177        with open(input, 'rb') as input_file:
178            content_len = len(input_file.read())
179
180            if content_len > partition.size:
181                raise Exception('Input file size exceeds partition size')
182
183        self._call_esptool(['write_flash', str(partition.offset), input] + self.esptool_write_args)
184
185
186def _write_partition(target, partition_id, input):
187    target.write_partition(partition_id, input)
188    partition = target.get_partition_info(partition_id)
189    status("Written contents of file '{}' at offset 0x{:x}".format(input, partition.offset))
190
191
192def _read_partition(target, partition_id, output):
193    target.read_partition(partition_id, output)
194    partition = target.get_partition_info(partition_id)
195    status("Read partition '{}' contents from device at offset 0x{:x} to file '{}'"
196           .format(partition.name, partition.offset, output))
197
198
199def _erase_partition(target, partition_id):
200    target.erase_partition(partition_id)
201    partition = target.get_partition_info(partition_id)
202    status("Erased partition '{}' at offset 0x{:x}".format(partition.name, partition.offset))
203
204
205def _get_partition_info(target, partition_id, info):
206    try:
207        partitions = target.get_partition_info(partition_id)
208        if not isinstance(partitions, list):
209            partitions = [partitions]
210    except Exception:
211        return
212
213    infos = []
214
215    try:
216        for p in partitions:
217            info_dict = {
218                'name': '{}'.format(p.name),
219                'type': '{}'.format(p.type),
220                'subtype': '{}'.format(p.subtype),
221                'offset': '0x{:x}'.format(p.offset),
222                'size': '0x{:x}'.format(p.size),
223                'encrypted': '{}'.format(p.encrypted)
224            }
225            for i in info:
226                infos += [info_dict[i]]
227    except KeyError:
228        raise RuntimeError('Request for unknown partition info {}'.format(i))
229
230    print(' '.join(infos))
231
232
233def main():
234    global quiet
235
236    parser = argparse.ArgumentParser('ESP-IDF Partitions Tool')
237
238    parser.add_argument('--quiet', '-q', help='suppress stderr messages', action='store_true')
239    parser.add_argument('--esptool-args', help='additional main arguments for esptool', nargs='+')
240    parser.add_argument('--esptool-write-args', help='additional subcommand arguments when writing to flash', nargs='+')
241    parser.add_argument('--esptool-read-args', help='additional subcommand arguments when reading flash', nargs='+')
242    parser.add_argument('--esptool-erase-args', help='additional subcommand arguments when erasing regions of flash', nargs='+')
243
244    # By default the device attached to the specified port is queried for the partition table. If a partition table file
245    # is specified, that is used instead.
246    parser.add_argument('--port', '-p', help='port where the target device of the command is connected to; the partition table is sourced from this device \
247                                            when the partition table file is not defined')
248    parser.add_argument('--baud', '-b', help='baudrate to use', type=int)
249
250    parser.add_argument('--partition-table-offset', '-o', help='offset to read the partition table from', type=str)
251    parser.add_argument('--partition-table-file', '-f', help='file (CSV/binary) to read the partition table from; \
252                                                            overrides device attached to specified port as the partition table source when defined')
253
254    partition_selection_parser = argparse.ArgumentParser(add_help=False)
255
256    # Specify what partition to perform the operation on. This can either be specified using the
257    # partition name or the first partition that matches the specified type/subtype
258    partition_selection_args = partition_selection_parser.add_mutually_exclusive_group()
259
260    partition_selection_args.add_argument('--partition-name', '-n', help='name of the partition')
261    partition_selection_args.add_argument('--partition-type', '-t', help='type of the partition')
262    partition_selection_args.add_argument('--partition-boot-default', '-d', help='select the default boot partition \
263                                           using the same fallback logic as the IDF bootloader', action='store_true')
264
265    partition_selection_parser.add_argument('--partition-subtype', '-s', help='subtype of the partition')
266
267    subparsers = parser.add_subparsers(dest='operation', help='run parttool -h for additional help')
268
269    # Specify the supported operations
270    read_part_subparser = subparsers.add_parser('read_partition', help='read partition from device and dump contents into a file',
271                                                parents=[partition_selection_parser])
272    read_part_subparser.add_argument('--output', help='file to dump the read partition contents to')
273
274    write_part_subparser = subparsers.add_parser('write_partition', help='write contents of a binary file to partition on device',
275                                                 parents=[partition_selection_parser])
276    write_part_subparser.add_argument('--input', help='file whose contents are to be written to the partition offset')
277
278    subparsers.add_parser('erase_partition', help='erase the contents of a partition on the device', parents=[partition_selection_parser])
279
280    print_partition_info_subparser = subparsers.add_parser('get_partition_info', help='get partition information', parents=[partition_selection_parser])
281    print_partition_info_subparser.add_argument('--info', help='type of partition information to get',
282                                                choices=['name', 'type', 'subtype', 'offset', 'size', 'encrypted'], default=['offset', 'size'], nargs='+')
283    print_partition_info_subparser.add_argument('--part_list', help='Get a list of partitions suitable for a given type', action='store_true')
284
285    args = parser.parse_args()
286    quiet = args.quiet
287
288    # No operation specified, display help and exit
289    if args.operation is None:
290        if not quiet:
291            parser.print_help()
292        sys.exit(1)
293
294    # Prepare the partition to perform operation on
295    if args.partition_name:
296        partition_id = PartitionName(args.partition_name)
297    elif args.partition_type:
298        if not args.partition_subtype:
299            raise RuntimeError('--partition-subtype should be defined when --partition-type is defined')
300        partition_id = PartitionType(args.partition_type, args.partition_subtype, getattr(args, 'part_list', None))
301    elif args.partition_boot_default:
302        partition_id = PARTITION_BOOT_DEFAULT
303    else:
304        raise RuntimeError('Partition to operate on should be defined using --partition-name OR \
305                            partition-type,--partition-subtype OR partition-boot-default')
306
307    # Prepare the device to perform operation on
308    target_args = {}
309
310    if args.port:
311        target_args['port'] = args.port
312
313    if args.baud:
314        target_args['baud'] = args.baud
315
316    if args.partition_table_file:
317        target_args['partition_table_file'] = args.partition_table_file
318
319    if args.partition_table_offset:
320        target_args['partition_table_offset'] = int(args.partition_table_offset, 0)
321
322    if args.esptool_args:
323        target_args['esptool_args'] = args.esptool_args
324
325    if args.esptool_write_args:
326        target_args['esptool_write_args'] = args.esptool_write_args
327
328    if args.esptool_read_args:
329        target_args['esptool_read_args'] = args.esptool_read_args
330
331    if args.esptool_erase_args:
332        target_args['esptool_erase_args'] = args.esptool_erase_args
333
334    target = ParttoolTarget(**target_args)
335
336    # Create the operation table and execute the operation
337    common_args = {'target':target, 'partition_id':partition_id}
338    parttool_ops = {
339        'erase_partition':(_erase_partition, []),
340        'read_partition':(_read_partition, ['output']),
341        'write_partition':(_write_partition, ['input']),
342        'get_partition_info':(_get_partition_info, ['info'])
343    }
344
345    (op, op_args) = parttool_ops[args.operation]
346
347    for op_arg in op_args:
348        common_args.update({op_arg:vars(args)[op_arg]})
349
350    if quiet:
351        # If exceptions occur, suppress and exit quietly
352        try:
353            op(**common_args)
354        except Exception:
355            sys.exit(2)
356    else:
357        try:
358            op(**common_args)
359        except gen.InputError as e:
360            print(e, file=sys.stderr)
361            sys.exit(2)
362
363
364if __name__ == '__main__':
365    main()
366