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