1#!/usr/bin/python
2# SPDX-License-Identifier: GPL-2.0+
3#
4# Copyright (C) 2016 Google, Inc
5# Written by Simon Glass <sjg@chromium.org>
6#
7
8from enum import IntEnum
9import struct
10import sys
11
12from dtoc import fdt_util
13import libfdt
14from libfdt import QUIET_NOTFOUND
15from patman import tools
16
17# This deals with a device tree, presenting it as an assortment of Node and
18# Prop objects, representing nodes and properties, respectively. This file
19# contains the base classes and defines the high-level API. You can use
20# FdtScan() as a convenience function to create and scan an Fdt.
21
22# This implementation uses a libfdt Python library to access the device tree,
23# so it is fairly efficient.
24
25# A list of types we support
26class Type(IntEnum):
27    (BYTE, INT, STRING, BOOL, INT64) = range(5)
28
29    def is_wider_than(self, other):
30        """Check if another type is 'wider' than this one
31
32        A wider type is one that holds more information than an earlier one,
33        similar to the concept of type-widening in C.
34
35        This uses a simple arithmetic comparison, since type values are in order
36        from narrowest (BYTE) to widest (INT64).
37
38        Args:
39            other: Other type to compare against
40
41        Return:
42            True if the other type is wider
43        """
44        return self.value > other.value
45
46def CheckErr(errnum, msg):
47    if errnum:
48        raise ValueError('Error %d: %s: %s' %
49            (errnum, libfdt.fdt_strerror(errnum), msg))
50
51
52def BytesToValue(data):
53    """Converts a string of bytes into a type and value
54
55    Args:
56        A bytes value (which on Python 2 is an alias for str)
57
58    Return:
59        A tuple:
60            Type of data
61            Data, either a single element or a list of elements. Each element
62            is one of:
63                Type.STRING: str/bytes value from the property
64                Type.INT: a byte-swapped integer stored as a 4-byte str/bytes
65                Type.BYTE: a byte stored as a single-byte str/bytes
66    """
67    data = bytes(data)
68    size = len(data)
69    strings = data.split(b'\0')
70    is_string = True
71    count = len(strings) - 1
72    if count > 0 and not len(strings[-1]):
73        for string in strings[:-1]:
74            if not string:
75                is_string = False
76                break
77            for ch in string:
78                if ch < 32 or ch > 127:
79                    is_string = False
80                    break
81    else:
82        is_string = False
83    if is_string:
84        if count == 1:
85            return Type.STRING, strings[0].decode()
86        else:
87            return Type.STRING, [s.decode() for s in strings[:-1]]
88    if size % 4:
89        if size == 1:
90            return Type.BYTE, chr(data[0])
91        else:
92            return Type.BYTE, [chr(ch) for ch in list(data)]
93    val = []
94    for i in range(0, size, 4):
95        val.append(data[i:i + 4])
96    if size == 4:
97        return Type.INT, val[0]
98    else:
99        return Type.INT, val
100
101
102class Prop:
103    """A device tree property
104
105    Properties:
106        name: Property name (as per the device tree)
107        value: Property value as a string of bytes, or a list of strings of
108            bytes
109        type: Value type
110    """
111    def __init__(self, node, offset, name, data):
112        self._node = node
113        self._offset = offset
114        self.name = name
115        self.value = None
116        self.bytes = bytes(data)
117        self.dirty = False
118        if not data:
119            self.type = Type.BOOL
120            self.value = True
121            return
122        self.type, self.value = BytesToValue(bytes(data))
123
124    def RefreshOffset(self, poffset):
125        self._offset = poffset
126
127    def Widen(self, newprop):
128        """Figure out which property type is more general
129
130        Given a current property and a new property, this function returns the
131        one that is less specific as to type. The less specific property will
132        be ble to represent the data in the more specific property. This is
133        used for things like:
134
135            node1 {
136                compatible = "fred";
137                value = <1>;
138            };
139            node1 {
140                compatible = "fred";
141                value = <1 2>;
142            };
143
144        He we want to use an int array for 'value'. The first property
145        suggests that a single int is enough, but the second one shows that
146        it is not. Calling this function with these two propertes would
147        update the current property to be like the second, since it is less
148        specific.
149        """
150        if self.type.is_wider_than(newprop.type):
151            if self.type == Type.INT and newprop.type == Type.BYTE:
152                if type(self.value) == list:
153                    new_value = []
154                    for val in self.value:
155                        new_value += [chr(by) for by in val]
156                else:
157                    new_value = [chr(by) for by in self.value]
158                self.value = new_value
159            self.type = newprop.type
160
161        if type(newprop.value) == list and type(self.value) != list:
162            self.value = [self.value]
163
164        if type(self.value) == list and len(newprop.value) > len(self.value):
165            val = self.GetEmpty(self.type)
166            while len(self.value) < len(newprop.value):
167                self.value.append(val)
168
169    @classmethod
170    def GetEmpty(self, type):
171        """Get an empty / zero value of the given type
172
173        Returns:
174            A single value of the given type
175        """
176        if type == Type.BYTE:
177            return chr(0)
178        elif type == Type.INT:
179            return struct.pack('>I', 0);
180        elif type == Type.STRING:
181            return ''
182        else:
183            return True
184
185    def GetOffset(self):
186        """Get the offset of a property
187
188        Returns:
189            The offset of the property (struct fdt_property) within the file
190        """
191        self._node._fdt.CheckCache()
192        return self._node._fdt.GetStructOffset(self._offset)
193
194    def SetInt(self, val):
195        """Set the integer value of the property
196
197        The device tree is marked dirty so that the value will be written to
198        the block on the next sync.
199
200        Args:
201            val: Integer value (32-bit, single cell)
202        """
203        self.bytes = struct.pack('>I', val);
204        self.value = self.bytes
205        self.type = Type.INT
206        self.dirty = True
207
208    def SetData(self, bytes):
209        """Set the value of a property as bytes
210
211        Args:
212            bytes: New property value to set
213        """
214        self.bytes = bytes
215        self.type, self.value = BytesToValue(bytes)
216        self.dirty = True
217
218    def Sync(self, auto_resize=False):
219        """Sync property changes back to the device tree
220
221        This updates the device tree blob with any changes to this property
222        since the last sync.
223
224        Args:
225            auto_resize: Resize the device tree automatically if it does not
226                have enough space for the update
227
228        Raises:
229            FdtException if auto_resize is False and there is not enough space
230        """
231        if self._offset is None or self.dirty:
232            node = self._node
233            fdt_obj = node._fdt._fdt_obj
234            if auto_resize:
235                while fdt_obj.setprop(node.Offset(), self.name, self.bytes,
236                                    (libfdt.NOSPACE,)) == -libfdt.NOSPACE:
237                    fdt_obj.resize(fdt_obj.totalsize() + 1024 +
238                                   len(self.bytes))
239                    fdt_obj.setprop(node.Offset(), self.name, self.bytes)
240            else:
241                fdt_obj.setprop(node.Offset(), self.name, self.bytes)
242
243
244class Node:
245    """A device tree node
246
247    Properties:
248        offset: Integer offset in the device tree
249        name: Device tree node tname
250        path: Full path to node, along with the node name itself
251        _fdt: Device tree object
252        subnodes: A list of subnodes for this node, each a Node object
253        props: A dict of properties for this node, each a Prop object.
254            Keyed by property name
255    """
256    def __init__(self, fdt, parent, offset, name, path):
257        self._fdt = fdt
258        self.parent = parent
259        self._offset = offset
260        self.name = name
261        self.path = path
262        self.subnodes = []
263        self.props = {}
264
265    def GetFdt(self):
266        """Get the Fdt object for this node
267
268        Returns:
269            Fdt object
270        """
271        return self._fdt
272
273    def FindNode(self, name):
274        """Find a node given its name
275
276        Args:
277            name: Node name to look for
278        Returns:
279            Node object if found, else None
280        """
281        for subnode in self.subnodes:
282            if subnode.name == name:
283                return subnode
284        return None
285
286    def Offset(self):
287        """Returns the offset of a node, after checking the cache
288
289        This should be used instead of self._offset directly, to ensure that
290        the cache does not contain invalid offsets.
291        """
292        self._fdt.CheckCache()
293        return self._offset
294
295    def Scan(self):
296        """Scan a node's properties and subnodes
297
298        This fills in the props and subnodes properties, recursively
299        searching into subnodes so that the entire tree is built.
300        """
301        fdt_obj = self._fdt._fdt_obj
302        self.props = self._fdt.GetProps(self)
303        phandle = fdt_obj.get_phandle(self.Offset())
304        if phandle:
305            self._fdt.phandle_to_node[phandle] = self
306
307        offset = fdt_obj.first_subnode(self.Offset(), QUIET_NOTFOUND)
308        while offset >= 0:
309            sep = '' if self.path[-1] == '/' else '/'
310            name = fdt_obj.get_name(offset)
311            path = self.path + sep + name
312            node = Node(self._fdt, self, offset, name, path)
313            self.subnodes.append(node)
314
315            node.Scan()
316            offset = fdt_obj.next_subnode(offset, QUIET_NOTFOUND)
317
318    def Refresh(self, my_offset):
319        """Fix up the _offset for each node, recursively
320
321        Note: This does not take account of property offsets - these will not
322        be updated.
323        """
324        fdt_obj = self._fdt._fdt_obj
325        if self._offset != my_offset:
326            self._offset = my_offset
327        offset = fdt_obj.first_subnode(self._offset, QUIET_NOTFOUND)
328        for subnode in self.subnodes:
329            if subnode.name != fdt_obj.get_name(offset):
330                raise ValueError('Internal error, node name mismatch %s != %s' %
331                                 (subnode.name, fdt_obj.get_name(offset)))
332            subnode.Refresh(offset)
333            offset = fdt_obj.next_subnode(offset, QUIET_NOTFOUND)
334        if offset != -libfdt.FDT_ERR_NOTFOUND:
335            raise ValueError('Internal error, offset == %d' % offset)
336
337        poffset = fdt_obj.first_property_offset(self._offset, QUIET_NOTFOUND)
338        while poffset >= 0:
339            p = fdt_obj.get_property_by_offset(poffset)
340            prop = self.props.get(p.name)
341            if not prop:
342                raise ValueError("Internal error, property '%s' missing, "
343                                 'offset %d' % (p.name, poffset))
344            prop.RefreshOffset(poffset)
345            poffset = fdt_obj.next_property_offset(poffset, QUIET_NOTFOUND)
346
347    def DeleteProp(self, prop_name):
348        """Delete a property of a node
349
350        The property is deleted and the offset cache is invalidated.
351
352        Args:
353            prop_name: Name of the property to delete
354        Raises:
355            ValueError if the property does not exist
356        """
357        CheckErr(self._fdt._fdt_obj.delprop(self.Offset(), prop_name),
358                 "Node '%s': delete property: '%s'" % (self.path, prop_name))
359        del self.props[prop_name]
360        self._fdt.Invalidate()
361
362    def AddZeroProp(self, prop_name):
363        """Add a new property to the device tree with an integer value of 0.
364
365        Args:
366            prop_name: Name of property
367        """
368        self.props[prop_name] = Prop(self, None, prop_name,
369                                     tools.GetBytes(0, 4))
370
371    def AddEmptyProp(self, prop_name, len):
372        """Add a property with a fixed data size, for filling in later
373
374        The device tree is marked dirty so that the value will be written to
375        the blob on the next sync.
376
377        Args:
378            prop_name: Name of property
379            len: Length of data in property
380        """
381        value = tools.GetBytes(0, len)
382        self.props[prop_name] = Prop(self, None, prop_name, value)
383
384    def _CheckProp(self, prop_name):
385        """Check if a property is present
386
387        Args:
388            prop_name: Name of property
389
390        Returns:
391            self
392
393        Raises:
394            ValueError if the property is missing
395        """
396        if prop_name not in self.props:
397            raise ValueError("Fdt '%s', node '%s': Missing property '%s'" %
398                             (self._fdt._fname, self.path, prop_name))
399        return self
400
401    def SetInt(self, prop_name, val):
402        """Update an integer property int the device tree.
403
404        This is not allowed to change the size of the FDT.
405
406        The device tree is marked dirty so that the value will be written to
407        the blob on the next sync.
408
409        Args:
410            prop_name: Name of property
411            val: Value to set
412        """
413        self._CheckProp(prop_name).props[prop_name].SetInt(val)
414
415    def SetData(self, prop_name, val):
416        """Set the data value of a property
417
418        The device tree is marked dirty so that the value will be written to
419        the blob on the next sync.
420
421        Args:
422            prop_name: Name of property to set
423            val: Data value to set
424        """
425        self._CheckProp(prop_name).props[prop_name].SetData(val)
426
427    def SetString(self, prop_name, val):
428        """Set the string value of a property
429
430        The device tree is marked dirty so that the value will be written to
431        the blob on the next sync.
432
433        Args:
434            prop_name: Name of property to set
435            val: String value to set (will be \0-terminated in DT)
436        """
437        if type(val) == str:
438            val = val.encode('utf-8')
439        self._CheckProp(prop_name).props[prop_name].SetData(val + b'\0')
440
441    def AddData(self, prop_name, val):
442        """Add a new property to a node
443
444        The device tree is marked dirty so that the value will be written to
445        the blob on the next sync.
446
447        Args:
448            prop_name: Name of property to add
449            val: Bytes value of property
450        """
451        self.props[prop_name] = Prop(self, None, prop_name, val)
452
453    def AddString(self, prop_name, val):
454        """Add a new string property to a node
455
456        The device tree is marked dirty so that the value will be written to
457        the blob on the next sync.
458
459        Args:
460            prop_name: Name of property to add
461            val: String value of property
462        """
463        val = bytes(val, 'utf-8')
464        self.AddData(prop_name, val + b'\0')
465
466    def AddInt(self, prop_name, val):
467        """Add a new integer property to a node
468
469        The device tree is marked dirty so that the value will be written to
470        the blob on the next sync.
471
472        Args:
473            prop_name: Name of property to add
474            val: Integer value of property
475        """
476        self.AddData(prop_name, struct.pack('>I', val))
477
478    def AddSubnode(self, name):
479        """Add a new subnode to the node
480
481        Args:
482            name: name of node to add
483
484        Returns:
485            New subnode that was created
486        """
487        path = self.path + '/' + name
488        subnode = Node(self._fdt, self, None, name, path)
489        self.subnodes.append(subnode)
490        return subnode
491
492    def Sync(self, auto_resize=False):
493        """Sync node changes back to the device tree
494
495        This updates the device tree blob with any changes to this node and its
496        subnodes since the last sync.
497
498        Args:
499            auto_resize: Resize the device tree automatically if it does not
500                have enough space for the update
501
502        Raises:
503            FdtException if auto_resize is False and there is not enough space
504        """
505        if self._offset is None:
506            # The subnode doesn't exist yet, so add it
507            fdt_obj = self._fdt._fdt_obj
508            if auto_resize:
509                while True:
510                    offset = fdt_obj.add_subnode(self.parent._offset, self.name,
511                                                (libfdt.NOSPACE,))
512                    if offset != -libfdt.NOSPACE:
513                        break
514                    fdt_obj.resize(fdt_obj.totalsize() + 1024)
515            else:
516                offset = fdt_obj.add_subnode(self.parent._offset, self.name)
517            self._offset = offset
518
519        # Sync subnodes in reverse so that we don't disturb node offsets for
520        # nodes that are earlier in the DT. This avoids an O(n^2) rescan of
521        # node offsets.
522        for node in reversed(self.subnodes):
523            node.Sync(auto_resize)
524
525        # Sync properties now, whose offsets should not have been disturbed.
526        # We do this after subnodes, since this disturbs the offsets of these
527        # properties. Note that new properties will have an offset of None here,
528        # which Python 3 cannot sort against int. So use a large value instead
529        # to ensure that the new properties are added first.
530        prop_list = sorted(self.props.values(),
531                           key=lambda prop: prop._offset or 1 << 31,
532                           reverse=True)
533        for prop in prop_list:
534            prop.Sync(auto_resize)
535
536
537class Fdt:
538    """Provides simple access to a flat device tree blob using libfdts.
539
540    Properties:
541      fname: Filename of fdt
542      _root: Root of device tree (a Node object)
543      name: Helpful name for this Fdt for the user (useful when creating the
544        DT from data rather than a file)
545    """
546    def __init__(self, fname):
547        self._fname = fname
548        self._cached_offsets = False
549        self.phandle_to_node = {}
550        self.name = ''
551        if self._fname:
552            self.name = self._fname
553            self._fname = fdt_util.EnsureCompiled(self._fname)
554
555            with open(self._fname, 'rb') as fd:
556                self._fdt_obj = libfdt.Fdt(fd.read())
557
558    @staticmethod
559    def FromData(data, name=''):
560        """Create a new Fdt object from the given data
561
562        Args:
563            data: Device-tree data blob
564            name: Helpful name for this Fdt for the user
565
566        Returns:
567            Fdt object containing the data
568        """
569        fdt = Fdt(None)
570        fdt._fdt_obj = libfdt.Fdt(bytes(data))
571        fdt.name = name
572        return fdt
573
574    def LookupPhandle(self, phandle):
575        """Look up a phandle
576
577        Args:
578            phandle: Phandle to look up (int)
579
580        Returns:
581            Node object the phandle points to
582        """
583        return self.phandle_to_node.get(phandle)
584
585    def Scan(self, root='/'):
586        """Scan a device tree, building up a tree of Node objects
587
588        This fills in the self._root property
589
590        Args:
591            root: Ignored
592
593        TODO(sjg@chromium.org): Implement the 'root' parameter
594        """
595        self._cached_offsets = True
596        self._root = self.Node(self, None, 0, '/', '/')
597        self._root.Scan()
598
599    def GetRoot(self):
600        """Get the root Node of the device tree
601
602        Returns:
603            The root Node object
604        """
605        return self._root
606
607    def GetNode(self, path):
608        """Look up a node from its path
609
610        Args:
611            path: Path to look up, e.g. '/microcode/update@0'
612        Returns:
613            Node object, or None if not found
614        """
615        node = self._root
616        parts = path.split('/')
617        if len(parts) < 2:
618            return None
619        if len(parts) == 2 and parts[1] == '':
620            return node
621        for part in parts[1:]:
622            node = node.FindNode(part)
623            if not node:
624                return None
625        return node
626
627    def Flush(self):
628        """Flush device tree changes back to the file
629
630        If the device tree has changed in memory, write it back to the file.
631        """
632        with open(self._fname, 'wb') as fd:
633            fd.write(self._fdt_obj.as_bytearray())
634
635    def Sync(self, auto_resize=False):
636        """Make sure any DT changes are written to the blob
637
638        Args:
639            auto_resize: Resize the device tree automatically if it does not
640                have enough space for the update
641
642        Raises:
643            FdtException if auto_resize is False and there is not enough space
644        """
645        self._root.Sync(auto_resize)
646        self.Invalidate()
647
648    def Pack(self):
649        """Pack the device tree down to its minimum size
650
651        When nodes and properties shrink or are deleted, wasted space can
652        build up in the device tree binary.
653        """
654        CheckErr(self._fdt_obj.pack(), 'pack')
655        self.Invalidate()
656
657    def GetContents(self):
658        """Get the contents of the FDT
659
660        Returns:
661            The FDT contents as a string of bytes
662        """
663        return bytes(self._fdt_obj.as_bytearray())
664
665    def GetFdtObj(self):
666        """Get the contents of the FDT
667
668        Returns:
669            The FDT contents as a libfdt.Fdt object
670        """
671        return self._fdt_obj
672
673    def GetProps(self, node):
674        """Get all properties from a node.
675
676        Args:
677            node: Full path to node name to look in.
678
679        Returns:
680            A dictionary containing all the properties, indexed by node name.
681            The entries are Prop objects.
682
683        Raises:
684            ValueError: if the node does not exist.
685        """
686        props_dict = {}
687        poffset = self._fdt_obj.first_property_offset(node._offset,
688                                                      QUIET_NOTFOUND)
689        while poffset >= 0:
690            p = self._fdt_obj.get_property_by_offset(poffset)
691            prop = Prop(node, poffset, p.name, p)
692            props_dict[prop.name] = prop
693
694            poffset = self._fdt_obj.next_property_offset(poffset,
695                                                         QUIET_NOTFOUND)
696        return props_dict
697
698    def Invalidate(self):
699        """Mark our offset cache as invalid"""
700        self._cached_offsets = False
701
702    def CheckCache(self):
703        """Refresh the offset cache if needed"""
704        if self._cached_offsets:
705            return
706        self.Refresh()
707        self._cached_offsets = True
708
709    def Refresh(self):
710        """Refresh the offset cache"""
711        self._root.Refresh(0)
712
713    def GetStructOffset(self, offset):
714        """Get the file offset of a given struct offset
715
716        Args:
717            offset: Offset within the 'struct' region of the device tree
718        Returns:
719            Position of @offset within the device tree binary
720        """
721        return self._fdt_obj.off_dt_struct() + offset
722
723    @classmethod
724    def Node(self, fdt, parent, offset, name, path):
725        """Create a new node
726
727        This is used by Fdt.Scan() to create a new node using the correct
728        class.
729
730        Args:
731            fdt: Fdt object
732            parent: Parent node, or None if this is the root node
733            offset: Offset of node
734            name: Node name
735            path: Full path to node
736        """
737        node = Node(fdt, parent, offset, name, path)
738        return node
739
740    def GetFilename(self):
741        """Get the filename of the device tree
742
743        Returns:
744            String filename
745        """
746        return self._fname
747
748def FdtScan(fname):
749    """Returns a new Fdt object"""
750    dtb = Fdt(fname)
751    dtb.Scan()
752    return dtb
753