1#! /usr/bin/env python
2# Copyright 2017, The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16from __future__ import print_function
17
18"""Tool for packing multiple DTB/DTBO files into a single image"""
19
20import argparse
21import os
22from array import array
23from collections import namedtuple
24import struct
25from sys import stdout
26import zlib
27
28class CompressionFormat(object):
29    """Enum representing DT compression format for a DT entry.
30    """
31    NO_COMPRESSION = 0x00
32    ZLIB_COMPRESSION = 0x01
33    GZIP_COMPRESSION = 0x02
34
35class DtEntry(object):
36    """Provides individual DT image file arguments to be added to a DTBO.
37
38    Attributes:
39        _REQUIRED_KEYS: 'keys' needed to be present in the dictionary passed to instantiate
40            an object of this class.
41        _COMPRESSION_FORMAT_MASK: Mask to retrieve compression info for DT entry from flags field
42            when a DTBO header of version 1 is used.
43    """
44    _COMPRESSION_FORMAT_MASK = 0x0f
45    REQUIRED_KEYS = ('dt_file', 'dt_size', 'dt_offset', 'id', 'rev', 'flags',
46                     'custom0', 'custom1', 'custom2')
47
48    @staticmethod
49    def __get_number_or_prop(arg):
50        """Converts string to integer or reads the property from DT image.
51
52        Args:
53            arg: String containing the argument provided on the command line.
54
55        Returns:
56            An integer property read from DT file or argument string
57            converted to integer
58        """
59
60        if not arg or arg[0] == '+' or arg[0] == '-':
61            raise ValueError('Invalid argument passed to DTImage')
62        if arg[0] == '/':
63            # TODO(b/XXX): Use pylibfdt to get property value from DT
64            raise ValueError('Invalid argument passed to DTImage')
65        else:
66            base = 10
67            if arg.startswith('0x') or arg.startswith('0X'):
68                base = 16
69            elif arg.startswith('0'):
70                base = 8
71            return int(arg, base)
72
73    def __init__(self, **kwargs):
74        """Constructor for DtEntry object.
75
76        Initializes attributes from dictionary object that contains
77        values keyed with names equivalent to the class's attributes.
78
79        Args:
80            kwargs: Dictionary object containing values to instantiate
81                class members with. Expected keys in dictionary are from
82                the tuple (_REQUIRED_KEYS)
83        """
84
85        missing_keys = set(self.REQUIRED_KEYS) - set(kwargs)
86        if missing_keys:
87            raise ValueError('Missing keys in DtEntry constructor: %r' %
88                             sorted(missing_keys))
89
90        self.__dt_file = kwargs['dt_file']
91        self.__dt_offset = kwargs['dt_offset']
92        self.__dt_size = kwargs['dt_size']
93        self.__id = self.__get_number_or_prop(kwargs['id'])
94        self.__rev = self.__get_number_or_prop(kwargs['rev'])
95        self.__flags = self.__get_number_or_prop(kwargs['flags'])
96        self.__custom0 = self.__get_number_or_prop(kwargs['custom0'])
97        self.__custom1 = self.__get_number_or_prop(kwargs['custom1'])
98        self.__custom2 = self.__get_number_or_prop(kwargs['custom2'])
99
100    def __str__(self):
101        sb = []
102        sb.append('{key:>20} = {value:d}'.format(key='dt_size',
103                                                 value=self.__dt_size))
104        sb.append('{key:>20} = {value:d}'.format(key='dt_offset',
105                                                 value=self.__dt_offset))
106        sb.append('{key:>20} = {value:08x}'.format(key='id',
107                                                   value=self.__id))
108        sb.append('{key:>20} = {value:08x}'.format(key='rev',
109                                                   value=self.__rev))
110        sb.append('{key:>20} = {value:08x}'.format(key='custom[0]',
111                                                   value=self.__flags))
112        sb.append('{key:>20} = {value:08x}'.format(key='custom[1]',
113                                                   value=self.__custom0))
114        sb.append('{key:>20} = {value:08x}'.format(key='custom[2]',
115                                                   value=self.__custom1))
116        sb.append('{key:>20} = {value:08x}'.format(key='custom[3]',
117                                                   value=self.__custom2))
118        return '\n'.join(sb)
119
120    def compression_info(self, version):
121        """CompressionFormat: compression format for DT image file.
122
123           Args:
124                version: Version of DTBO header, compression is only
125                         supported from version 1.
126        """
127        if version is 0:
128            return CompressionFormat.NO_COMPRESSION
129        return self.flags & self._COMPRESSION_FORMAT_MASK
130
131    @property
132    def dt_file(self):
133        """file: File handle to the DT image file."""
134        return self.__dt_file
135
136    @property
137    def size(self):
138        """int: size in bytes of the DT image file."""
139        return self.__dt_size
140
141    @size.setter
142    def size(self, value):
143        self.__dt_size = value
144
145    @property
146    def dt_offset(self):
147        """int: offset in DTBO file for this DT image."""
148        return self.__dt_offset
149
150    @dt_offset.setter
151    def dt_offset(self, value):
152        self.__dt_offset = value
153
154    @property
155    def image_id(self):
156        """int: DT entry _id for this DT image."""
157        return self.__id
158
159    @property
160    def rev(self):
161        """int: DT entry _rev for this DT image."""
162        return self.__rev
163
164    @property
165    def flags(self):
166        """int: DT entry _flags for this DT image."""
167        return self.__flags
168
169    @property
170    def custom0(self):
171        """int: DT entry _custom0 for this DT image."""
172        return self.__custom0
173
174    @property
175    def custom1(self):
176        """int: DT entry _custom1 for this DT image."""
177        return self.__custom1
178
179    @property
180    def custom2(self):
181        """int: DT entry custom2 for this DT image."""
182        return self.__custom2
183
184
185class Dtbo(object):
186    """
187    Provides parser, reader, writer for dumping and creating Device Tree Blob
188    Overlay (DTBO) images.
189
190    Attributes:
191        _DTBO_MAGIC: Device tree table header magic.
192        _ACPIO_MAGIC: Advanced Configuration and Power Interface table header
193                      magic.
194        _DT_TABLE_HEADER_SIZE: Size of Device tree table header.
195        _DT_TABLE_HEADER_INTS: Number of integers in DT table header.
196        _DT_ENTRY_HEADER_SIZE: Size of Device tree entry header within a DTBO.
197        _DT_ENTRY_HEADER_INTS: Number of integers in DT entry header.
198        _GZIP_COMPRESSION_WBITS: Argument 'wbits' for gzip compression
199        _ZLIB_DECOMPRESSION_WBITS: Argument 'wbits' for zlib/gzip compression
200    """
201
202    _DTBO_MAGIC = 0xd7b7ab1e
203    _ACPIO_MAGIC = 0x41435049
204    _DT_TABLE_HEADER_SIZE = struct.calcsize('>8I')
205    _DT_TABLE_HEADER_INTS = 8
206    _DT_ENTRY_HEADER_SIZE = struct.calcsize('>8I')
207    _DT_ENTRY_HEADER_INTS = 8
208    _GZIP_COMPRESSION_WBITS = 31
209    _ZLIB_DECOMPRESSION_WBITS = 47
210
211    def _update_dt_table_header(self):
212        """Converts header entries into binary data for DTBO header.
213
214        Packs the current Device tree table header attribute values in
215        metadata buffer.
216        """
217        struct.pack_into('>8I', self.__metadata, 0, self.magic,
218                         self.total_size, self.header_size,
219                         self.dt_entry_size, self.dt_entry_count,
220                         self.dt_entries_offset, self.page_size,
221                         self.version)
222
223    def _update_dt_entry_header(self, dt_entry, metadata_offset):
224        """Converts each DT entry header entry into binary data for DTBO file.
225
226        Packs the current device tree table entry attribute into
227        metadata buffer as device tree entry header.
228
229        Args:
230            dt_entry: DtEntry object for the header to be packed.
231            metadata_offset: Offset into metadata buffer to begin writing.
232            dtbo_offset: Offset where the DT image file for this dt_entry can
233                be found in the resulting DTBO image.
234        """
235        struct.pack_into('>8I', self.__metadata, metadata_offset, dt_entry.size,
236                         dt_entry.dt_offset, dt_entry.image_id, dt_entry.rev,
237                         dt_entry.flags, dt_entry.custom0, dt_entry.custom1,
238                         dt_entry.custom2)
239
240    def _update_metadata(self):
241        """Updates the DTBO metadata.
242
243        Initialize the internal metadata buffer and fill it with all Device
244        Tree table entries and update the DTBO header.
245        """
246
247        self.__metadata = array('c', ' ' * self.__metadata_size)
248        metadata_offset = self.header_size
249        for dt_entry in self.__dt_entries:
250            self._update_dt_entry_header(dt_entry, metadata_offset)
251            metadata_offset += self.dt_entry_size
252        self._update_dt_table_header()
253
254    def _read_dtbo_header(self, buf):
255        """Reads DTBO file header into metadata buffer.
256
257        Unpack and read the DTBO table header from given buffer. The
258        buffer size must exactly be equal to _DT_TABLE_HEADER_SIZE.
259
260        Args:
261            buf: Bytebuffer read directly from the file of size
262                _DT_TABLE_HEADER_SIZE.
263        """
264        (self.magic, self.total_size, self.header_size,
265         self.dt_entry_size, self.dt_entry_count, self.dt_entries_offset,
266         self.page_size, self.version) = struct.unpack_from('>8I', buf, 0)
267
268        # verify the header
269        if self.magic != self._DTBO_MAGIC and self.magic != self._ACPIO_MAGIC:
270            raise ValueError('Invalid magic number 0x%x in DTBO/ACPIO file' %
271                             (self.magic))
272
273        if self.header_size != self._DT_TABLE_HEADER_SIZE:
274            raise ValueError('Invalid header size (%d) in DTBO/ACPIO file' %
275                             (self.header_size))
276
277        if self.dt_entry_size != self._DT_ENTRY_HEADER_SIZE:
278            raise ValueError('Invalid DT entry header size (%d) in DTBO/ACPIO file' %
279                             (self.dt_entry_size))
280
281    def _read_dt_entries_from_metadata(self):
282        """Reads individual DT entry headers from metadata buffer.
283
284        Unpack and read the DTBO DT entry headers from the internal buffer.
285        The buffer size must exactly be equal to _DT_TABLE_HEADER_SIZE +
286        (_DT_ENTRY_HEADER_SIZE * dt_entry_count). The method raises exception
287        if DT entries have already been set for this object.
288        """
289
290        if self.__dt_entries:
291            raise ValueError('DTBO DT entries can be added only once')
292
293        offset = self.dt_entries_offset / 4
294        params = {}
295        params['dt_file'] = None
296        for i in range(0, self.dt_entry_count):
297            dt_table_entry = self.__metadata[offset:offset + self._DT_ENTRY_HEADER_INTS]
298            params['dt_size'] = dt_table_entry[0]
299            params['dt_offset'] = dt_table_entry[1]
300            for j in range(2, self._DT_ENTRY_HEADER_INTS):
301                params[DtEntry.REQUIRED_KEYS[j + 1]] = str(dt_table_entry[j])
302            dt_entry = DtEntry(**params)
303            self.__dt_entries.append(dt_entry)
304            offset += self._DT_ENTRY_HEADER_INTS
305
306    def _read_dtbo_image(self):
307        """Parse the input file and instantiate this object."""
308
309        # First check if we have enough to read the header
310        file_size = os.fstat(self.__file.fileno()).st_size
311        if file_size < self._DT_TABLE_HEADER_SIZE:
312            raise ValueError('Invalid DTBO file')
313
314        self.__file.seek(0)
315        buf = self.__file.read(self._DT_TABLE_HEADER_SIZE)
316        self._read_dtbo_header(buf)
317
318        self.__metadata_size = (self.header_size +
319                                self.dt_entry_count * self.dt_entry_size)
320        if file_size < self.__metadata_size:
321            raise ValueError('Invalid or truncated DTBO file of size %d expected %d' %
322                             file_size, self.__metadata_size)
323
324        num_ints = (self._DT_TABLE_HEADER_INTS +
325                    self.dt_entry_count * self._DT_ENTRY_HEADER_INTS)
326        if self.dt_entries_offset > self._DT_TABLE_HEADER_SIZE:
327            num_ints += (self.dt_entries_offset - self._DT_TABLE_HEADER_SIZE) / 4
328        format_str = '>' + str(num_ints) + 'I'
329        self.__file.seek(0)
330        self.__metadata = struct.unpack(format_str,
331                                        self.__file.read(self.__metadata_size))
332        self._read_dt_entries_from_metadata()
333
334    def _find_dt_entry_with_same_file(self, dt_entry):
335        """Finds DT Entry that has identical backing DT file.
336
337        Args:
338            dt_entry: DtEntry object whose 'dtfile' we find for existence in the
339                current 'dt_entries'.
340        Returns:
341            If a match by file path is found, the corresponding DtEntry object
342            from internal list is returned. If not, 'None' is returned.
343        """
344
345        dt_entry_path = os.path.realpath(dt_entry.dt_file.name)
346        for entry in self.__dt_entries:
347            entry_path = os.path.realpath(entry.dt_file.name)
348            if entry_path == dt_entry_path:
349                return entry
350        return None
351
352    def __init__(self, file_handle, dt_type='dtb', page_size=None, version=0):
353        """Constructor for Dtbo Object
354
355        Args:
356            file_handle: The Dtbo File handle corresponding to this object.
357                The file handle can be used to write to (in case of 'create')
358                or read from (in case of 'dump')
359        """
360
361        self.__file = file_handle
362        self.__dt_entries = []
363        self.__metadata = None
364        self.__metadata_size = 0
365
366        # if page_size is given, assume the object is being instantiated to
367        # create a DTBO file
368        if page_size:
369            if dt_type == 'acpi':
370                self.magic = self._ACPIO_MAGIC
371            else:
372                self.magic = self._DTBO_MAGIC
373            self.total_size = self._DT_TABLE_HEADER_SIZE
374            self.header_size = self._DT_TABLE_HEADER_SIZE
375            self.dt_entry_size = self._DT_ENTRY_HEADER_SIZE
376            self.dt_entry_count = 0
377            self.dt_entries_offset = self._DT_TABLE_HEADER_SIZE
378            self.page_size = page_size
379            self.version = version
380            self.__metadata_size = self._DT_TABLE_HEADER_SIZE
381        else:
382            self._read_dtbo_image()
383
384    def __str__(self):
385        sb = []
386        sb.append('dt_table_header:')
387        _keys = ('magic', 'total_size', 'header_size', 'dt_entry_size',
388                 'dt_entry_count', 'dt_entries_offset', 'page_size', 'version')
389        for key in _keys:
390            if key == 'magic':
391                sb.append('{key:>20} = {value:08x}'.format(key=key,
392                                                           value=self.__dict__[key]))
393            else:
394                sb.append('{key:>20} = {value:d}'.format(key=key,
395                                                         value=self.__dict__[key]))
396        count = 0
397        for dt_entry in self.__dt_entries:
398            sb.append('dt_table_entry[{0:d}]:'.format(count))
399            sb.append(str(dt_entry))
400            count = count + 1
401        return '\n'.join(sb)
402
403    @property
404    def dt_entries(self):
405        """Returns a list of DtEntry objects found in DTBO file."""
406        return self.__dt_entries
407
408    def compress_dt_entry(self, compression_format, dt_entry_file):
409        """Compresses a DT entry.
410
411        Args:
412            compression_format: Compression format for DT Entry
413            dt_entry_file: File handle to read DT entry from.
414
415        Returns:
416            Compressed DT entry and its length.
417
418        Raises:
419            ValueError if unrecognized compression format is found.
420        """
421        compress_zlib = zlib.compressobj()  #  zlib
422        compress_gzip = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
423                                         zlib.DEFLATED, self._GZIP_COMPRESSION_WBITS)  #  gzip
424        compression_obj_dict = {
425            CompressionFormat.NO_COMPRESSION: None,
426            CompressionFormat.ZLIB_COMPRESSION: compress_zlib,
427            CompressionFormat.GZIP_COMPRESSION: compress_gzip,
428        }
429
430        if compression_format not in compression_obj_dict:
431            ValueError("Bad compression format %d" % compression_format)
432
433        if compression_format is CompressionFormat.NO_COMPRESSION:
434            dt_entry = dt_entry_file.read()
435        else:
436            compression_object = compression_obj_dict[compression_format]
437            dt_entry_file.seek(0)
438            dt_entry = compression_object.compress(dt_entry_file.read())
439            dt_entry += compression_object.flush()
440        return dt_entry, len(dt_entry)
441
442    def add_dt_entries(self, dt_entries):
443        """Adds DT image files to the DTBO object.
444
445        Adds a list of Dtentry Objects to the DTBO image. The changes are not
446        committed to the output file until commit() is called.
447
448        Args:
449            dt_entries: List of DtEntry object to be added.
450
451        Returns:
452            A buffer containing all DT entries.
453
454        Raises:
455            ValueError: if the list of DT entries is empty or if a list of DT entries
456                has already been added to the DTBO.
457        """
458        if not dt_entries:
459            raise ValueError('Attempted to add empty list of DT entries')
460
461        if self.__dt_entries:
462            raise ValueError('DTBO DT entries can be added only once')
463
464        dt_entry_count = len(dt_entries)
465        dt_offset = (self.header_size +
466                     dt_entry_count * self.dt_entry_size)
467
468        dt_entry_buf = ""
469        for dt_entry in dt_entries:
470            if not isinstance(dt_entry, DtEntry):
471                raise ValueError('Adding invalid DT entry object to DTBO')
472            entry = self._find_dt_entry_with_same_file(dt_entry)
473            dt_entry_compression_info = dt_entry.compression_info(self.version)
474            if entry and (entry.compression_info(self.version)
475                          == dt_entry_compression_info):
476                dt_entry.dt_offset = entry.dt_offset
477                dt_entry.size = entry.size
478            else:
479                dt_entry.dt_offset = dt_offset
480                compressed_entry, dt_entry.size = self.compress_dt_entry(dt_entry_compression_info,
481                                                                         dt_entry.dt_file)
482                dt_entry_buf += compressed_entry
483                dt_offset += dt_entry.size
484                self.total_size += dt_entry.size
485            self.__dt_entries.append(dt_entry)
486            self.dt_entry_count += 1
487            self.__metadata_size += self.dt_entry_size
488            self.total_size += self.dt_entry_size
489
490        return dt_entry_buf
491
492    def extract_dt_file(self, idx, fout, decompress):
493        """Extract DT Image files embedded in the DTBO file.
494
495        Extracts Device Tree blob image file at given index into a file handle.
496
497        Args:
498            idx: Index of the DT entry in the DTBO file.
499            fout: File handle where the DTB at index idx to be extracted into.
500            decompress: If a DT entry is compressed, decompress it before writing
501                it to the file handle.
502
503        Raises:
504            ValueError: if invalid DT entry index or compression format is detected.
505        """
506        if idx > self.dt_entry_count:
507            raise ValueError('Invalid index %d of DtEntry' % idx)
508
509        size = self.dt_entries[idx].size
510        offset = self.dt_entries[idx].dt_offset
511        self.__file.seek(offset, 0)
512        fout.seek(0)
513        compression_format = self.dt_entries[idx].compression_info(self.version)
514        if decompress and compression_format:
515            if (compression_format == CompressionFormat.ZLIB_COMPRESSION or
516                compression_format == CompressionFormat.GZIP_COMPRESSION):
517                fout.write(zlib.decompress(self.__file.read(size), self._ZLIB_DECOMPRESSION_WBITS))
518            else:
519                raise ValueError("Unknown compression format detected")
520        else:
521            fout.write(self.__file.read(size))
522
523    def commit(self, dt_entry_buf):
524        """Write out staged changes to the DTBO object to create a DTBO file.
525
526        Writes a fully instantiated Dtbo Object into the output file using the
527        file handle present in '_file'. No checks are performed on the object
528        except for existence of output file handle on the object before writing
529        out the file.
530
531        Args:
532            dt_entry_buf: Buffer containing all DT entries.
533        """
534        if not self.__file:
535            raise ValueError('No file given to write to.')
536
537        if not self.__dt_entries:
538            raise ValueError('No DT image files to embed into DTBO image given.')
539
540        self._update_metadata()
541
542        self.__file.seek(0)
543        self.__file.write(self.__metadata)
544        self.__file.write(dt_entry_buf)
545        self.__file.flush()
546
547
548def parse_dt_entry(global_args, arglist):
549    """Parse arguments for single DT entry file.
550
551    Parses command line arguments for single DT image file while
552    creating a Device tree blob overlay (DTBO).
553
554    Args:
555        global_args: Dtbo object containing global default values
556            for DtEntry attributes.
557        arglist: Command line argument list for this DtEntry.
558
559    Returns:
560        A Namespace object containing all values to instantiate DtEntry object.
561    """
562
563    parser = argparse.ArgumentParser(add_help=False)
564    parser.add_argument('dt_file', nargs='?',
565                        type=argparse.FileType('rb'),
566                        default=None)
567    parser.add_argument('--id', type=str, dest='id', action='store',
568                        default=global_args.global_id)
569    parser.add_argument('--rev', type=str, dest='rev',
570                        action='store', default=global_args.global_rev)
571    parser.add_argument('--flags', type=str, dest='flags',
572                        action='store',
573                        default=global_args.global_flags)
574    parser.add_argument('--custom0', type=str, dest='custom0',
575                        action='store',
576                        default=global_args.global_custom0)
577    parser.add_argument('--custom1', type=str, dest='custom1',
578                        action='store',
579                        default=global_args.global_custom1)
580    parser.add_argument('--custom2', type=str, dest='custom2',
581                        action='store',
582                        default=global_args.global_custom2)
583    return parser.parse_args(arglist)
584
585
586def parse_dt_entries(global_args, arg_list):
587    """Parse all DT entries from command line.
588
589    Parse all DT image files and their corresponding attribute from
590    command line
591
592    Args:
593        global_args: Argument containing default global values for _id,
594            _rev and customX.
595        arg_list: The remainder of the command line after global options
596            DTBO creation have been parsed.
597
598    Returns:
599        A List of DtEntry objects created after parsing the command line
600        given in argument.
601    """
602    dt_entries = []
603    img_file_idx = []
604    idx = 0
605    # find all positional arguments (i.e. DT image file paths)
606    for arg in arg_list:
607        if not arg.startswith("--"):
608            img_file_idx.append(idx)
609        idx = idx + 1
610
611    if not img_file_idx:
612        raise ValueError('Input DT images must be provided')
613
614    total_images = len(img_file_idx)
615    for idx in xrange(total_images):
616        start_idx = img_file_idx[idx]
617        if idx == total_images - 1:
618            argv = arg_list[start_idx:]
619        else:
620            end_idx = img_file_idx[idx + 1]
621            argv = arg_list[start_idx:end_idx]
622        args = parse_dt_entry(global_args, argv)
623        params = vars(args)
624        params['dt_offset'] = 0
625        params['dt_size'] = os.fstat(params['dt_file'].fileno()).st_size
626        dt_entries.append(DtEntry(**params))
627
628    return dt_entries
629
630def parse_config_option(line, is_global, dt_keys, global_key_types):
631    """Parses a single line from the configuration file.
632
633    Args:
634        line: String containing the key=value line from the file.
635        is_global: Boolean indicating if we should parse global or DT entry
636            specific option.
637        dt_keys: Tuple containing all valid DT entry and global option strings
638            in configuration file.
639        global_key_types: A dict of global options and their corresponding types. It
640            contains all exclusive valid global option strings in configuration
641            file that are not repeated in dt entry options.
642
643    Returns:
644        Returns a tuple for parsed key and value for the option. Also, checks
645        the key to make sure its valid.
646    """
647
648    if line.find('=') == -1:
649        raise ValueError('Invalid line (%s) in configuration file' % line)
650
651    key, value = (x.strip() for x in line.split('='))
652    if is_global and key in global_key_types:
653        if global_key_types[key] is int:
654            value = int(value)
655    elif key not in dt_keys:
656        raise ValueError('Invalid option (%s) in configuration file' % key)
657
658    return key, value
659
660def parse_config_file(fin, dt_keys, global_key_types):
661    """Parses the configuration file for creating DTBO image.
662
663    Args:
664        fin: File handle for configuration file
665        is_global: Boolean indicating if we should parse global or DT entry
666            specific option.
667        dt_keys: Tuple containing all valid DT entry and global option strings
668            in configuration file.
669        global_key_types: A dict of global options and their corresponding types. It
670            contains all exclusive valid global option strings in configuration
671            file that are not repeated in dt entry options.
672
673    Returns:
674        global_args, dt_args: Tuple of a dictionary with global arguments
675        and a list of dictionaries for all DT entry specific arguments the
676        following format.
677            global_args:
678                {'id' : <value>, 'rev' : <value> ...}
679            dt_args:
680                [{'filename' : 'dt_file_name', 'id' : <value>,
681                 'rev' : <value> ...},
682                 {'filename' : 'dt_file_name2', 'id' : <value2>,
683                  'rev' : <value2> ...}, ...
684                ]
685    """
686
687    # set all global defaults
688    global_args = dict((k, '0') for k in dt_keys)
689    global_args['dt_type'] = 'dtb'
690    global_args['page_size'] = 2048
691    global_args['version'] = 0
692
693    dt_args = []
694    found_dt_entry = False
695    count = -1
696    for line in fin:
697        line = line.rstrip()
698        if line.lstrip().startswith('#'):
699            continue
700        comment_idx = line.find('#')
701        line = line if comment_idx == -1 else line[0:comment_idx]
702        if not line or line.isspace():
703            continue
704        if line.startswith((' ', '\t')) and not found_dt_entry:
705            # This is a global argument
706            key, value = parse_config_option(line, True, dt_keys, global_key_types)
707            global_args[key] = value
708        elif line.find('=') != -1:
709            key, value = parse_config_option(line, False, dt_keys, global_key_types)
710            dt_args[-1][key] = value
711        else:
712            found_dt_entry = True
713            count += 1
714            dt_args.append({})
715            dt_args[-1]['filename'] = line.strip()
716    return global_args, dt_args
717
718def parse_create_args(arg_list):
719    """Parse command line arguments for 'create' sub-command.
720
721    Args:
722        arg_list: All command line arguments except the outfile file name.
723
724    Returns:
725        The list of remainder of the command line arguments after parsing
726        for 'create'.
727    """
728
729    image_arg_index = 0
730    for arg in arg_list:
731        if not arg.startswith("--"):
732            break
733        image_arg_index = image_arg_index + 1
734
735    argv = arg_list[0:image_arg_index]
736    remainder = arg_list[image_arg_index:]
737    parser = argparse.ArgumentParser(prog='create', add_help=False)
738    parser.add_argument('--dt_type', type=str, dest='dt_type',
739                        action='store', default='dtb')
740    parser.add_argument('--page_size', type=int, dest='page_size',
741                        action='store', default=2048)
742    parser.add_argument('--version', type=int, dest='version',
743                        action='store', default=0)
744    parser.add_argument('--id', type=str, dest='global_id',
745                        action='store', default='0')
746    parser.add_argument('--rev', type=str, dest='global_rev',
747                        action='store', default='0')
748    parser.add_argument('--flags', type=str, dest='global_flags',
749                        action='store', default='0')
750    parser.add_argument('--custom0', type=str, dest='global_custom0',
751                        action='store', default='0')
752    parser.add_argument('--custom1', type=str, dest='global_custom1',
753                        action='store', default='0')
754    parser.add_argument('--custom2', type=str, dest='global_custom2',
755                        action='store', default='0')
756    args = parser.parse_args(argv)
757    return args, remainder
758
759def parse_dump_cmd_args(arglist):
760    """Parse command line arguments for 'dump' sub-command.
761
762    Args:
763        arglist: List of all command line arguments including the outfile
764            file name if exists.
765
766    Returns:
767        A namespace object of parsed arguments.
768    """
769
770    parser = argparse.ArgumentParser(prog='dump')
771    parser.add_argument('--output', '-o', nargs='?',
772                        type=argparse.FileType('wb'),
773                        dest='outfile',
774                        default=stdout)
775    parser.add_argument('--dtb', '-b', nargs='?', type=str,
776                        dest='dtfilename')
777    parser.add_argument('--decompress', action='store_true', dest='decompress')
778    return parser.parse_args(arglist)
779
780def parse_config_create_cmd_args(arglist):
781    """Parse command line arguments for 'cfg_create subcommand.
782
783    Args:
784        arglist: A list of all command line arguments including the
785            mandatory input configuration file name.
786
787    Returns:
788        A Namespace object of parsed arguments.
789    """
790    parser = argparse.ArgumentParser(prog='cfg_create')
791    parser.add_argument('conf_file', nargs='?',
792                        type=argparse.FileType('rb'),
793                        default=None)
794    cwd = os.getcwd()
795    parser.add_argument('--dtb-dir', '-d', nargs='?', type=str,
796                        dest='dtbdir', default=cwd)
797    return parser.parse_args(arglist)
798
799def create_dtbo_image(fout, argv):
800    """Create Device Tree Blob Overlay image using provided arguments.
801
802    Args:
803        fout: Output file handle to write to.
804        argv: list of command line arguments.
805    """
806
807    global_args, remainder = parse_create_args(argv)
808    if not remainder:
809        raise ValueError('List of dtimages to add to DTBO not provided')
810    dt_entries = parse_dt_entries(global_args, remainder)
811    dtbo = Dtbo(fout, global_args.dt_type, global_args.page_size, global_args.version)
812    dt_entry_buf = dtbo.add_dt_entries(dt_entries)
813    dtbo.commit(dt_entry_buf)
814    fout.close()
815
816def dump_dtbo_image(fin, argv):
817    """Dump DTBO file.
818
819    Dump Device Tree Blob Overlay metadata as output and the device
820    tree image files embedded in the DTBO image into file(s) provided
821    as arguments
822
823    Args:
824        fin: Input DTBO image files.
825        argv: list of command line arguments.
826    """
827    dtbo = Dtbo(fin)
828    args = parse_dump_cmd_args(argv)
829    if args.dtfilename:
830        num_entries = len(dtbo.dt_entries)
831        for idx in range(0, num_entries):
832            with open(args.dtfilename + '.{:d}'.format(idx), 'wb') as fout:
833                dtbo.extract_dt_file(idx, fout, args.decompress)
834    args.outfile.write(str(dtbo) + '\n')
835    args.outfile.close()
836
837def create_dtbo_image_from_config(fout, argv):
838    """Create DTBO file from a configuration file.
839
840    Args:
841        fout: Output file handle to write to.
842        argv: list of command line arguments.
843    """
844    args = parse_config_create_cmd_args(argv)
845    if not args.conf_file:
846        raise ValueError('Configuration file must be provided')
847
848    _DT_KEYS = ('id', 'rev', 'flags', 'custom0', 'custom1', 'custom2')
849    _GLOBAL_KEY_TYPES = {'dt_type': str, 'page_size': int, 'version': int}
850
851    global_args, dt_args = parse_config_file(args.conf_file,
852                                             _DT_KEYS, _GLOBAL_KEY_TYPES)
853    params = {}
854    dt_entries = []
855    for dt_arg in dt_args:
856        filepath = args.dtbdir + os.sep + dt_arg['filename']
857        params['dt_file'] = open(filepath, 'rb')
858        params['dt_offset'] = 0
859        params['dt_size'] = os.fstat(params['dt_file'].fileno()).st_size
860        for key in _DT_KEYS:
861            if key not in dt_arg:
862                params[key] = global_args[key]
863            else:
864                params[key] = dt_arg[key]
865        dt_entries.append(DtEntry(**params))
866
867    # Create and write DTBO file
868    dtbo = Dtbo(fout, global_args['dt_type'], global_args['page_size'], global_args['version'])
869    dt_entry_buf = dtbo.add_dt_entries(dt_entries)
870    dtbo.commit(dt_entry_buf)
871    fout.close()
872
873def print_default_usage(progname):
874    """Prints program's default help string.
875
876    Args:
877        progname: This program's name.
878    """
879    sb = []
880    sb.append('  ' + progname + ' help all')
881    sb.append('  ' + progname + ' help <command>\n')
882    sb.append('    commands:')
883    sb.append('      help, dump, create, cfg_create')
884    print('\n'.join(sb))
885
886def print_dump_usage(progname):
887    """Prints usage for 'dump' sub-command.
888
889    Args:
890        progname: This program's name.
891    """
892    sb = []
893    sb.append('  ' + progname + ' dump <image_file> (<option>...)\n')
894    sb.append('    options:')
895    sb.append('      -o, --output <filename>  Output file name.')
896    sb.append('                               Default is output to stdout.')
897    sb.append('      -b, --dtb <filename>     Dump dtb/dtbo files from image.')
898    sb.append('                               Will output to <filename>.0, <filename>.1, etc.')
899    print('\n'.join(sb))
900
901def print_create_usage(progname):
902    """Prints usage for 'create' subcommand.
903
904    Args:
905        progname: This program's name.
906    """
907    sb = []
908    sb.append('  ' + progname + ' create <image_file> (<global_option>...) (<dtb_file> (<entry_option>...) ...)\n')
909    sb.append('    global_options:')
910    sb.append('      --dt_type=<type>         Device Tree Type (dtb|acpi). Default: dtb')
911    sb.append('      --page_size=<number>     Page size. Default: 2048')
912    sb.append('      --version=<number>       DTBO/ACPIO version. Default: 0')
913    sb.append('      --id=<number>       The default value to set property id in dt_table_entry. Default: 0')
914    sb.append('      --rev=<number>')
915    sb.append('      --flags=<number>')
916    sb.append('      --custom0=<number>')
917    sb.append('      --custom1=<number>')
918    sb.append('      --custom2=<number>\n')
919
920    sb.append('      The value could be a number or a DT node path.')
921    sb.append('      <number> could be a 32-bits digit or hex value, ex. 68000, 0x6800.')
922    sb.append('      <path> format is <full_node_path>:<property_name>, ex. /board/:id,')
923    sb.append('      will read the value in given FTB file with the path.')
924    print('\n'.join(sb))
925
926def print_cfg_create_usage(progname):
927    """Prints usage for 'cfg_create' sub-command.
928
929    Args:
930        progname: This program's name.
931    """
932    sb = []
933    sb.append('  ' + progname + ' cfg_create <image_file> <config_file> (<option>...)\n')
934    sb.append('    options:')
935    sb.append('      -d, --dtb-dir <dir>      The path to load dtb files.')
936    sb.append('                               Default is load from the current path.')
937    print('\n'.join(sb))
938
939def print_usage(cmd, _):
940    """Prints usage for this program.
941
942    Args:
943        cmd: The string sub-command for which help (usage) is requested.
944    """
945    prog_name = os.path.basename(__file__)
946    if not cmd:
947        print_default_usage(prog_name)
948        return
949
950    HelpCommand = namedtuple('HelpCommand', 'help_cmd, help_func')
951    help_commands = (HelpCommand('dump', print_dump_usage),
952                     HelpCommand('create', print_create_usage),
953                     HelpCommand('cfg_create', print_cfg_create_usage),
954                     )
955
956    if cmd == 'all':
957        print_default_usage(prog_name)
958
959    for help_cmd, help_func in help_commands:
960        if cmd == 'all' or cmd == help_cmd:
961            help_func(prog_name)
962            if cmd != 'all':
963                return
964
965    print('Unsupported help command: %s' % cmd, end='\n\n')
966    print_default_usage(prog_name)
967    return
968
969def main():
970    """Main entry point for mkdtboimg."""
971
972    parser = argparse.ArgumentParser(prog='mkdtboimg.py')
973
974    subparser = parser.add_subparsers(title='subcommand',
975                                      description='Valid subcommands')
976
977    create_parser = subparser.add_parser('create', add_help=False)
978    create_parser.add_argument('argfile', nargs='?',
979                               action='store', help='Output File',
980                               type=argparse.FileType('wb'))
981    create_parser.set_defaults(func=create_dtbo_image)
982
983    config_parser = subparser.add_parser('cfg_create', add_help=False)
984    config_parser.add_argument('argfile', nargs='?',
985                               action='store',
986                               type=argparse.FileType('wb'))
987    config_parser.set_defaults(func=create_dtbo_image_from_config)
988
989    dump_parser = subparser.add_parser('dump', add_help=False)
990    dump_parser.add_argument('argfile', nargs='?',
991                             action='store',
992                             type=argparse.FileType('rb'))
993    dump_parser.set_defaults(func=dump_dtbo_image)
994
995    help_parser = subparser.add_parser('help', add_help=False)
996    help_parser.add_argument('argfile', nargs='?', action='store')
997    help_parser.set_defaults(func=print_usage)
998
999    (subcmd, subcmd_args) = parser.parse_known_args()
1000    subcmd.func(subcmd.argfile, subcmd_args)
1001
1002if __name__ == '__main__':
1003    main()
1004