1#!/usr/bin/env python
2#
3# ESP32 partition table generation tool
4#
5# Converts partition tables to/from CSV and binary formats.
6#
7# See https://docs.espressif.com/projects/esp-idf/en/latest/api-guides/partition-tables.html
8# for explanation of partition table structure and uses.
9#
10# Copyright 2015-2016 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 binascii
27import errno
28import hashlib
29import os
30import re
31import struct
32import sys
33
34MAX_PARTITION_LENGTH = 0xC00   # 3K for partition data (96 entries) leaves 1K in a 4K sector for signature
35MD5_PARTITION_BEGIN = b'\xEB\xEB' + b'\xFF' * 14  # The first 2 bytes are like magic numbers for MD5 sum
36PARTITION_TABLE_SIZE  = 0x1000  # Size of partition table
37
38MIN_PARTITION_SUBTYPE_APP_OTA = 0x10
39NUM_PARTITION_SUBTYPE_APP_OTA = 16
40
41__version__ = '1.2'
42
43APP_TYPE = 0x00
44DATA_TYPE = 0x01
45
46TYPES = {
47    'app': APP_TYPE,
48    'data': DATA_TYPE,
49}
50
51
52def get_ptype_as_int(ptype):
53    """ Convert a string which might be numeric or the name of a partition type to an integer """
54    try:
55        return TYPES[ptype]
56    except KeyError:
57        try:
58            return int(ptype, 0)
59        except TypeError:
60            return ptype
61
62
63# Keep this map in sync with esp_partition_subtype_t enum in esp_partition.h
64SUBTYPES = {
65    APP_TYPE: {
66        'factory': 0x00,
67        'test': 0x20,
68    },
69    DATA_TYPE: {
70        'ota': 0x00,
71        'phy': 0x01,
72        'nvs': 0x02,
73        'coredump': 0x03,
74        'nvs_keys': 0x04,
75        'efuse': 0x05,
76        'undefined': 0x06,
77        'esphttpd': 0x80,
78        'fat': 0x81,
79        'spiffs': 0x82,
80    },
81}
82
83
84def get_subtype_as_int(ptype, subtype):
85    """ Convert a string which might be numeric or the name of a partition subtype to an integer """
86    try:
87        return SUBTYPES[get_ptype_as_int(ptype)][subtype]
88    except KeyError:
89        try:
90            return int(subtype, 0)
91        except TypeError:
92            return subtype
93
94
95ALIGNMENT = {
96    APP_TYPE: 0x10000,
97    DATA_TYPE: 0x4,
98}
99
100
101STRICT_DATA_ALIGNMENT = 0x1000
102
103
104def get_alignment_for_type(ptype):
105    return ALIGNMENT.get(ptype, ALIGNMENT[DATA_TYPE])
106
107
108quiet = False
109md5sum = True
110secure = False
111offset_part_table = 0
112
113
114def status(msg):
115    """ Print status message to stderr """
116    if not quiet:
117        critical(msg)
118
119
120def critical(msg):
121    """ Print critical message to stderr """
122    sys.stderr.write(msg)
123    sys.stderr.write('\n')
124
125
126class PartitionTable(list):
127    def __init__(self):
128        super(PartitionTable, self).__init__(self)
129
130    @classmethod
131    def from_file(cls, f):
132        data = f.read()
133        data_is_binary = data[0:2] == PartitionDefinition.MAGIC_BYTES
134        if data_is_binary:
135            status('Parsing binary partition input...')
136            return cls.from_binary(data), True
137
138        data = data.decode()
139        status('Parsing CSV input...')
140        return cls.from_csv(data), False
141
142    @classmethod
143    def from_csv(cls, csv_contents):
144        res = PartitionTable()
145        lines = csv_contents.splitlines()
146
147        def expand_vars(f):
148            f = os.path.expandvars(f)
149            m = re.match(r'(?<!\\)\$([A-Za-z_][A-Za-z0-9_]*)', f)
150            if m:
151                raise InputError("unknown variable '%s'" % m.group(1))
152            return f
153
154        for line_no in range(len(lines)):
155            line = expand_vars(lines[line_no]).strip()
156            if line.startswith('#') or len(line) == 0:
157                continue
158            try:
159                res.append(PartitionDefinition.from_csv(line, line_no + 1))
160            except InputError as err:
161                raise InputError('Error at line %d: %s' % (line_no + 1, err))
162            except Exception:
163                critical('Unexpected error parsing CSV line %d: %s' % (line_no + 1, line))
164                raise
165
166        # fix up missing offsets & negative sizes
167        last_end = offset_part_table + PARTITION_TABLE_SIZE  # first offset after partition table
168        for e in res:
169            if e.offset is not None and e.offset < last_end:
170                if e == res[0]:
171                    raise InputError('CSV Error: First partition offset 0x%x overlaps end of partition table 0x%x'
172                                     % (e.offset, last_end))
173                else:
174                    raise InputError('CSV Error: Partitions overlap. Partition at line %d sets offset 0x%x. Previous partition ends 0x%x'
175                                     % (e.line_no, e.offset, last_end))
176            if e.offset is None:
177                pad_to = get_alignment_for_type(e.type)
178                if last_end % pad_to != 0:
179                    last_end += pad_to - (last_end % pad_to)
180                e.offset = last_end
181            if e.size < 0:
182                e.size = -e.size - e.offset
183            last_end = e.offset + e.size
184
185        return res
186
187    def __getitem__(self, item):
188        """ Allow partition table access via name as well as by
189        numeric index. """
190        if isinstance(item, str):
191            for x in self:
192                if x.name == item:
193                    return x
194            raise ValueError("No partition entry named '%s'" % item)
195        else:
196            return super(PartitionTable, self).__getitem__(item)
197
198    def find_by_type(self, ptype, subtype):
199        """ Return a partition by type & subtype, returns
200        None if not found """
201        # convert ptype & subtypes names (if supplied this way) to integer values
202        ptype = get_ptype_as_int(ptype)
203        subtype = get_subtype_as_int(ptype, subtype)
204
205        for p in self:
206            if p.type == ptype and p.subtype == subtype:
207                yield p
208        return
209
210    def find_by_name(self, name):
211        for p in self:
212            if p.name == name:
213                return p
214        return None
215
216    def verify(self):
217        # verify each partition individually
218        for p in self:
219            p.verify()
220
221        # check on duplicate name
222        names = [p.name for p in self]
223        duplicates = set(n for n in names if names.count(n) > 1)
224
225        # print sorted duplicate partitions by name
226        if len(duplicates) != 0:
227            critical('A list of partitions that have the same name:')
228            for p in sorted(self, key=lambda x:x.name):
229                if len(duplicates.intersection([p.name])) != 0:
230                    critical('%s' % (p.to_csv()))
231            raise InputError('Partition names must be unique')
232
233        # check for overlaps
234        last = None
235        for p in sorted(self, key=lambda x:x.offset):
236            if p.offset < offset_part_table + PARTITION_TABLE_SIZE:
237                raise InputError('Partition offset 0x%x is below 0x%x' % (p.offset, offset_part_table + PARTITION_TABLE_SIZE))
238            if last is not None and p.offset < last.offset + last.size:
239                raise InputError('Partition at 0x%x overlaps 0x%x-0x%x' % (p.offset, last.offset, last.offset + last.size - 1))
240            last = p
241
242        # check that otadata should be unique
243        otadata_duplicates = [p for p in self if p.type == TYPES['data'] and p.subtype == SUBTYPES[DATA_TYPE]['ota']]
244        if len(otadata_duplicates) > 1:
245            for p in otadata_duplicates:
246                critical('%s' % (p.to_csv()))
247            raise InputError('Found multiple otadata partitions. Only one partition can be defined with type="data"(1) and subtype="ota"(0).')
248
249        if len(otadata_duplicates) == 1 and otadata_duplicates[0].size != 0x2000:
250            p = otadata_duplicates[0]
251            critical('%s' % (p.to_csv()))
252            raise InputError('otadata partition must have size = 0x2000')
253
254    def flash_size(self):
255        """ Return the size that partitions will occupy in flash
256            (ie the offset the last partition ends at)
257        """
258        try:
259            last = sorted(self, reverse=True)[0]
260        except IndexError:
261            return 0  # empty table!
262        return last.offset + last.size
263
264    @classmethod
265    def from_binary(cls, b):
266        md5 = hashlib.md5()
267        result = cls()
268        for o in range(0,len(b),32):
269            data = b[o:o + 32]
270            if len(data) != 32:
271                raise InputError('Partition table length must be a multiple of 32 bytes')
272            if data == b'\xFF' * 32:
273                return result  # got end marker
274            if md5sum and data[:2] == MD5_PARTITION_BEGIN[:2]:  # check only the magic number part
275                if data[16:] == md5.digest():
276                    continue  # the next iteration will check for the end marker
277                else:
278                    raise InputError("MD5 checksums don't match! (computed: 0x%s, parsed: 0x%s)" % (md5.hexdigest(), binascii.hexlify(data[16:])))
279            else:
280                md5.update(data)
281            result.append(PartitionDefinition.from_binary(data))
282        raise InputError('Partition table is missing an end-of-table marker')
283
284    def to_binary(self):
285        result = b''.join(e.to_binary() for e in self)
286        if md5sum:
287            result += MD5_PARTITION_BEGIN + hashlib.md5(result).digest()
288        if len(result) >= MAX_PARTITION_LENGTH:
289            raise InputError('Binary partition table length (%d) longer than max' % len(result))
290        result += b'\xFF' * (MAX_PARTITION_LENGTH - len(result))  # pad the sector, for signing
291        return result
292
293    def to_csv(self, simple_formatting=False):
294        rows = ['# ESP-IDF Partition Table',
295                '# Name, Type, SubType, Offset, Size, Flags']
296        rows += [x.to_csv(simple_formatting) for x in self]
297        return '\n'.join(rows) + '\n'
298
299
300class PartitionDefinition(object):
301    MAGIC_BYTES = b'\xAA\x50'
302
303    # dictionary maps flag name (as used in CSV flags list, property name)
304    # to bit set in flags words in binary format
305    FLAGS = {
306        'encrypted': 0
307    }
308
309    # add subtypes for the 16 OTA slot values ("ota_XX, etc.")
310    for ota_slot in range(NUM_PARTITION_SUBTYPE_APP_OTA):
311        SUBTYPES[TYPES['app']]['ota_%d' % ota_slot] = MIN_PARTITION_SUBTYPE_APP_OTA + ota_slot
312
313    def __init__(self):
314        self.name = ''
315        self.type = None
316        self.subtype = None
317        self.offset = None
318        self.size = None
319        self.encrypted = False
320
321    @classmethod
322    def from_csv(cls, line, line_no):
323        """ Parse a line from the CSV """
324        line_w_defaults = line + ',,,,'  # lazy way to support default fields
325        fields = [f.strip() for f in line_w_defaults.split(',')]
326
327        res = PartitionDefinition()
328        res.line_no = line_no
329        res.name = fields[0]
330        res.type = res.parse_type(fields[1])
331        res.subtype = res.parse_subtype(fields[2])
332        res.offset = res.parse_address(fields[3])
333        res.size = res.parse_address(fields[4])
334        if res.size is None:
335            raise InputError("Size field can't be empty")
336
337        flags = fields[5].split(':')
338        for flag in flags:
339            if flag in cls.FLAGS:
340                setattr(res, flag, True)
341            elif len(flag) > 0:
342                raise InputError("CSV flag column contains unknown flag '%s'" % (flag))
343
344        return res
345
346    def __eq__(self, other):
347        return self.name == other.name and self.type == other.type \
348            and self.subtype == other.subtype and self.offset == other.offset \
349            and self.size == other.size
350
351    def __repr__(self):
352        def maybe_hex(x):
353            return '0x%x' % x if x is not None else 'None'
354        return "PartitionDefinition('%s', 0x%x, 0x%x, %s, %s)" % (self.name, self.type, self.subtype or 0,
355                                                                  maybe_hex(self.offset), maybe_hex(self.size))
356
357    def __str__(self):
358        return "Part '%s' %d/%d @ 0x%x size 0x%x" % (self.name, self.type, self.subtype, self.offset or -1, self.size or -1)
359
360    def __cmp__(self, other):
361        return self.offset - other.offset
362
363    def __lt__(self, other):
364        return self.offset < other.offset
365
366    def __gt__(self, other):
367        return self.offset > other.offset
368
369    def __le__(self, other):
370        return self.offset <= other.offset
371
372    def __ge__(self, other):
373        return self.offset >= other.offset
374
375    def parse_type(self, strval):
376        if strval == '':
377            raise InputError("Field 'type' can't be left empty.")
378        return parse_int(strval, TYPES)
379
380    def parse_subtype(self, strval):
381        if strval == '':
382            if self.type == TYPES['app']:
383                raise InputError('App partition cannot have an empty subtype')
384            return SUBTYPES[DATA_TYPE]['undefined']
385        return parse_int(strval, SUBTYPES.get(self.type, {}))
386
387    def parse_address(self, strval):
388        if strval == '':
389            return None  # PartitionTable will fill in default
390        return parse_int(strval)
391
392    def verify(self):
393        if self.type is None:
394            raise ValidationError(self, 'Type field is not set')
395        if self.subtype is None:
396            raise ValidationError(self, 'Subtype field is not set')
397        if self.offset is None:
398            raise ValidationError(self, 'Offset field is not set')
399        align = get_alignment_for_type(self.type)
400        if self.offset % align:
401            raise ValidationError(self, 'Offset 0x%x is not aligned to 0x%x' % (self.offset, align))
402        # The alignment requirement for non-app partition is 4 bytes, but it should be 4 kB.
403        # Print a warning for now, make it an error in IDF 5.0 (IDF-3742).
404        if self.type != APP_TYPE and self.offset % STRICT_DATA_ALIGNMENT:
405            critical('WARNING: Partition %s not aligned to 0x%x.'
406                     'This is deprecated and will be considered an error in the future release.' % (self.name, STRICT_DATA_ALIGNMENT))
407        if self.size % align and secure and self.type == APP_TYPE:
408            raise ValidationError(self, 'Size 0x%x is not aligned to 0x%x' % (self.size, align))
409        if self.size is None:
410            raise ValidationError(self, 'Size field is not set')
411
412        if self.name in TYPES and TYPES.get(self.name, '') != self.type:
413            critical("WARNING: Partition has name '%s' which is a partition type, but does not match this partition's "
414                     'type (0x%x). Mistake in partition table?' % (self.name, self.type))
415        all_subtype_names = []
416        for names in (t.keys() for t in SUBTYPES.values()):
417            all_subtype_names += names
418        if self.name in all_subtype_names and SUBTYPES.get(self.type, {}).get(self.name, '') != self.subtype:
419            critical("WARNING: Partition has name '%s' which is a partition subtype, but this partition has "
420                     'non-matching type 0x%x and subtype 0x%x. Mistake in partition table?' % (self.name, self.type, self.subtype))
421
422    STRUCT_FORMAT = b'<2sBBLL16sL'
423
424    @classmethod
425    def from_binary(cls, b):
426        if len(b) != 32:
427            raise InputError('Partition definition length must be exactly 32 bytes. Got %d bytes.' % len(b))
428        res = cls()
429        (magic, res.type, res.subtype, res.offset,
430         res.size, res.name, flags) = struct.unpack(cls.STRUCT_FORMAT, b)
431        if b'\x00' in res.name:  # strip null byte padding from name string
432            res.name = res.name[:res.name.index(b'\x00')]
433        res.name = res.name.decode()
434        if magic != cls.MAGIC_BYTES:
435            raise InputError('Invalid magic bytes (%r) for partition definition' % magic)
436        for flag,bit in cls.FLAGS.items():
437            if flags & (1 << bit):
438                setattr(res, flag, True)
439                flags &= ~(1 << bit)
440        if flags != 0:
441            critical('WARNING: Partition definition had unknown flag(s) 0x%08x. Newer binary format?' % flags)
442        return res
443
444    def get_flags_list(self):
445        return [flag for flag in self.FLAGS.keys() if getattr(self, flag)]
446
447    def to_binary(self):
448        flags = sum((1 << self.FLAGS[flag]) for flag in self.get_flags_list())
449        return struct.pack(self.STRUCT_FORMAT,
450                           self.MAGIC_BYTES,
451                           self.type, self.subtype,
452                           self.offset, self.size,
453                           self.name.encode(),
454                           flags)
455
456    def to_csv(self, simple_formatting=False):
457        def addr_format(a, include_sizes):
458            if not simple_formatting and include_sizes:
459                for (val, suffix) in [(0x100000, 'M'), (0x400, 'K')]:
460                    if a % val == 0:
461                        return '%d%s' % (a // val, suffix)
462            return '0x%x' % a
463
464        def lookup_keyword(t, keywords):
465            for k,v in keywords.items():
466                if simple_formatting is False and t == v:
467                    return k
468            return '%d' % t
469
470        def generate_text_flags():
471            """ colon-delimited list of flags """
472            return ':'.join(self.get_flags_list())
473
474        return ','.join([self.name,
475                         lookup_keyword(self.type, TYPES),
476                         lookup_keyword(self.subtype, SUBTYPES.get(self.type, {})),
477                         addr_format(self.offset, False),
478                         addr_format(self.size, True),
479                         generate_text_flags()])
480
481
482def parse_int(v, keywords={}):
483    """Generic parser for integer fields - int(x,0) with provision for
484    k/m/K/M suffixes and 'keyword' value lookup.
485    """
486    try:
487        for letter, multiplier in [('k', 1024), ('m', 1024 * 1024)]:
488            if v.lower().endswith(letter):
489                return parse_int(v[:-1], keywords) * multiplier
490        return int(v, 0)
491    except ValueError:
492        if len(keywords) == 0:
493            raise InputError('Invalid field value %s' % v)
494        try:
495            return keywords[v.lower()]
496        except KeyError:
497            raise InputError("Value '%s' is not valid. Known keywords: %s" % (v, ', '.join(keywords)))
498
499
500def main():
501    global quiet
502    global md5sum
503    global offset_part_table
504    global secure
505    parser = argparse.ArgumentParser(description='ESP32 partition table utility')
506
507    parser.add_argument('--flash-size', help='Optional flash size limit, checks partition table fits in flash',
508                        nargs='?', choices=['1MB', '2MB', '4MB', '8MB', '16MB'])
509    parser.add_argument('--disable-md5sum', help='Disable md5 checksum for the partition table', default=False, action='store_true')
510    parser.add_argument('--no-verify', help="Don't verify partition table fields", action='store_true')
511    parser.add_argument('--verify', '-v', help='Verify partition table fields (deprecated, this behaviour is '
512                                               'enabled by default and this flag does nothing.', action='store_true')
513    parser.add_argument('--quiet', '-q', help="Don't print non-critical status messages to stderr", action='store_true')
514    parser.add_argument('--offset', '-o', help='Set offset partition table', default='0x8000')
515    parser.add_argument('--secure', help='Require app partitions to be suitable for secure boot', action='store_true')
516    parser.add_argument('input', help='Path to CSV or binary file to parse.', type=argparse.FileType('rb'))
517    parser.add_argument('output', help='Path to output converted binary or CSV file. Will use stdout if omitted.',
518                        nargs='?', default='-')
519
520    args = parser.parse_args()
521
522    quiet = args.quiet
523    md5sum = not args.disable_md5sum
524    secure = args.secure
525    offset_part_table = int(args.offset, 0)
526    table, input_is_binary = PartitionTable.from_file(args.input)
527
528    if not args.no_verify:
529        status('Verifying table...')
530        table.verify()
531
532    if args.flash_size:
533        size_mb = int(args.flash_size.replace('MB', ''))
534        size = size_mb * 1024 * 1024  # flash memory uses honest megabytes!
535        table_size = table.flash_size()
536        if size < table_size:
537            raise InputError("Partitions defined in '%s' occupy %.1fMB of flash (%d bytes) which does not fit in configured "
538                             "flash size %dMB. Change the flash size in menuconfig under the 'Serial Flasher Config' menu." %
539                             (args.input.name, table_size / 1024.0 / 1024.0, table_size, size_mb))
540
541    # Make sure that the output directory is created
542    output_dir = os.path.abspath(os.path.dirname(args.output))
543
544    if not os.path.exists(output_dir):
545        try:
546            os.makedirs(output_dir)
547        except OSError as exc:
548            if exc.errno != errno.EEXIST:
549                raise
550
551    if input_is_binary:
552        output = table.to_csv()
553        with sys.stdout if args.output == '-' else open(args.output, 'w') as f:
554            f.write(output)
555    else:
556        output = table.to_binary()
557        try:
558            stdout_binary = sys.stdout.buffer  # Python 3
559        except AttributeError:
560            stdout_binary = sys.stdout
561        with stdout_binary if args.output == '-' else open(args.output, 'wb') as f:
562            f.write(output)
563
564
565class InputError(RuntimeError):
566    def __init__(self, e):
567        super(InputError, self).__init__(e)
568
569
570class ValidationError(InputError):
571    def __init__(self, partition, message):
572        super(ValidationError, self).__init__(
573            'Partition %s invalid: %s' % (partition.name, message))
574
575
576if __name__ == '__main__':
577    try:
578        main()
579    except InputError as e:
580        print(e, file=sys.stderr)
581        sys.exit(2)
582