1#!/usr/bin/python
2# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""A module to support automated testing of ChromeOS firmware.
7
8Utilizes services provided by saft_flashrom_util.py read/write the
9flashrom chip and to parse the flash rom image.
10
11See docstring for FlashromHandler class below.
12"""
13
14import hashlib
15import os
16import struct
17
18class FvSection(object):
19    """An object to hold information about a firmware section.
20
21    This includes file names for the signature header and the body, and the
22    version number.
23    """
24
25    def __init__(self, sig_name, body_name):
26        self._sig_name = sig_name
27        self._body_name = body_name
28        self._version = -1  # Is not set on construction.
29        self._flags = 0  # Is not set on construction.
30        self._sha = None  # Is not set on construction.
31        self._sig_sha = None # Is not set on construction.
32        self._datakey_version = -1 # Is not set on construction.
33        self._kernel_subkey_version = -1 # Is not set on construction.
34
35    def names(self):
36        return (self._sig_name, self._body_name)
37
38    def get_sig_name(self):
39        return self._sig_name
40
41    def get_body_name(self):
42        return self._body_name
43
44    def get_version(self):
45        return self._version
46
47    def get_flags(self):
48        return self._flags
49
50    def get_sha(self):
51        return self._sha
52
53    def get_sig_sha(self):
54        return self._sig_sha
55
56    def get_datakey_version(self):
57        return self._datakey_version
58
59    def get_kernel_subkey_version(self):
60        return self._kernel_subkey_version
61
62    def set_version(self, version):
63        self._version = version
64
65    def set_flags(self, flags):
66        self._flags = flags
67
68    def set_sha(self, sha):
69        self._sha = sha
70
71    def set_sig_sha(self, sha):
72        self._sig_sha = sha
73
74    def set_datakey_version(self, version):
75        self._datakey_version = version
76
77    def set_kernel_subkey_version(self, version):
78        self._kernel_subkey_version = version
79
80class FlashromHandlerError(Exception):
81    pass
82
83
84class FlashromHandler(object):
85    """An object to provide logical services for automated flashrom testing."""
86
87    DELTA = 1  # value to add to a byte to corrupt a section contents
88
89    # File in the state directory to store public root key.
90    PUB_KEY_FILE_NAME = 'root.pubkey'
91    FW_KEYBLOCK_FILE_NAME = 'firmware.keyblock'
92    FW_PRIV_DATA_KEY_FILE_NAME = 'firmware_data_key.vbprivk'
93    KERNEL_SUBKEY_FILE_NAME = 'kernel_subkey.vbpubk'
94
95    def __init__(self):
96    # make sure it does not accidentally overwrite the image.
97        self.fum = None
98        self.os_if = None
99        self.image = ''
100        self.pub_key_file = ''
101
102    def init(self, flashrom_util_module,
103             os_if,
104             pub_key_file=None,
105             dev_key_path='./',
106             target='bios'):
107        """Flashrom handler initializer.
108
109        Args:
110          flashrom_util_module - a module providing flashrom access utilities.
111          os_if - a module providing interface to OS services
112          pub_key_file - a string, name of the file contaning a public key to
113                         use for verifying both existing and new firmware.
114        """
115        if target == 'bios':
116            self.fum = flashrom_util_module.flashrom_util(
117                    os_if, target_is_ec=False)
118            self.fv_sections = {
119                'a': FvSection('VBOOTA', 'FVMAIN'),
120                'b': FvSection('VBOOTB', 'FVMAINB'),
121                'ec_a': FvSection(None, 'ECMAINA'),
122                'ec_b': FvSection(None, 'ECMAINB'),
123                }
124        elif target == 'ec':
125            self.fum = flashrom_util_module.flashrom_util(
126                    os_if, target_is_ec=True)
127            self.fv_sections = {
128                'rw': FvSection(None, 'EC_RW'),
129                }
130        else:
131            raise FlashromHandlerError("Invalid target.")
132        self.os_if = os_if
133        self.pub_key_file = pub_key_file
134        self.dev_key_path = dev_key_path
135
136    def new_image(self, image_file=None):
137        """Parse the full flashrom image and store sections into files.
138
139        Args:
140          image_file - a string, the name of the file contaning full ChromeOS
141                       flashrom image. If not passed in or empty - the actual
142                       flashrom is read and its contents are saved into a
143                       temporary file which is used instead.
144
145        The input file is parsed and the sections of importance (as defined in
146        self.fv_sections) are saved in separate files in the state directory
147        as defined in the os_if object.
148        """
149
150        if image_file:
151            self.image = open(image_file, 'rb').read()
152            self.fum.set_firmware_layout(image_file)
153        else:
154            self.image = self.fum.read_whole()
155
156        for section in self.fv_sections.itervalues():
157            for subsection_name in section.names():
158                if not subsection_name:
159                    continue
160                blob = self.fum.get_section(self.image, subsection_name)
161                if blob:
162                    f = open(self.os_if.state_dir_file(subsection_name),
163                             'wb')
164                    f.write(blob)
165                    f.close()
166
167            blob = self.fum.get_section(self.image, section.get_body_name())
168            if blob:
169                s = hashlib.sha1()
170                s.update(blob)
171                section.set_sha(s.hexdigest())
172
173            # If there is no "sig" subsection, skip reading version and flags.
174            if not section.get_sig_name():
175                continue
176
177            # Now determine this section's version number.
178            vb_section = self.fum.get_section(
179                self.image, section.get_sig_name())
180
181            section.set_version(self.os_if.retrieve_body_version(vb_section))
182            section.set_flags(self.os_if.retrieve_preamble_flags(vb_section))
183            section.set_datakey_version(
184                self.os_if.retrieve_datakey_version(vb_section))
185            section.set_kernel_subkey_version(
186                self.os_if.retrieve_kernel_subkey_version(vb_section))
187
188            s = hashlib.sha1()
189            s.update(self.fum.get_section(self.image, section.get_sig_name()))
190            section.set_sig_sha(s.hexdigest())
191
192        if not self.pub_key_file:
193            self._retrieve_pub_key()
194
195    def _retrieve_pub_key(self):
196        """Retrieve root public key from the firmware GBB section."""
197
198        gbb_header_format = '<4s20s2I'
199        pubk_header_format = '<2Q'
200
201        gbb_section = self.fum.get_section(self.image, 'FV_GBB')
202
203        # do some sanity checks
204        try:
205            sig, _, rootk_offs, rootk_size = struct.unpack_from(
206                gbb_header_format, gbb_section)
207        except struct.error, e:
208            raise FlashromHandlerError(e)
209
210        if sig != '$GBB' or (rootk_offs + rootk_size) > len(gbb_section):
211            raise FlashromHandlerError('Bad gbb header')
212
213        key_body_offset, key_body_size = struct.unpack_from(
214            pubk_header_format, gbb_section, rootk_offs)
215
216        # Generally speaking the offset field can be anything, but in case of
217        # GBB section the key is stored as a standalone entity, so the offset
218        # of the key body is expected to be equal to the key header size of
219        # 0x20.
220        # Should this convention change, the check below would fail, which
221        # would be a good prompt for revisiting this test's behavior and
222        # algorithms.
223        if key_body_offset != 0x20 or key_body_size > rootk_size:
224            raise FlashromHandlerError('Bad public key format')
225
226        # All checks passed, let's store the key in a file.
227        self.pub_key_file = self.os_if.state_dir_file(self.PUB_KEY_FILE_NAME)
228        keyf = open(self.pub_key_file, 'w')
229        key = gbb_section[
230            rootk_offs:rootk_offs + key_body_offset + key_body_size]
231        keyf.write(key)
232        keyf.close()
233
234    def verify_image(self):
235        """Confirm the image's validity.
236
237        Using the file supplied to init() as the public key container verify
238        the two sections' (FirmwareA and FirmwareB) integrity. The contents of
239        the sections is taken from the files created by new_image()
240
241        In case there is an integrity error raises FlashromHandlerError
242        exception with the appropriate error message text.
243        """
244
245        for section in self.fv_sections.itervalues():
246            if section.get_sig_name():
247                cmd = 'vbutil_firmware --verify %s --signpubkey %s  --fv %s' % (
248                    self.os_if.state_dir_file(section.get_sig_name()),
249                    self.pub_key_file,
250                    self.os_if.state_dir_file(section.get_body_name()))
251                self.os_if.run_shell_command(cmd)
252
253    def _modify_section(self, section, delta, body_or_sig=False,
254                        corrupt_all=False):
255        """Modify a firmware section inside the image, either body or signature.
256
257        If corrupt_all is set, the passed in delta is added to all bytes in the
258        section. Otherwise, the delta is added to the value located at 2% offset
259        into the section blob, either body or signature.
260
261        Calling this function again for the same section the complimentary
262        delta value would restore the section contents.
263        """
264
265        if not self.image:
266            raise FlashromHandlerError(
267                'Attempt at using an uninitialized object')
268        if section not in self.fv_sections:
269            raise FlashromHandlerError('Unknown FW section %s'
270                                       % section)
271
272        # Get the appropriate section of the image.
273        if body_or_sig:
274            subsection_name = self.fv_sections[section].get_body_name()
275        else:
276            subsection_name = self.fv_sections[section].get_sig_name()
277        blob = self.fum.get_section(self.image, subsection_name)
278
279        # Modify the byte in it within 2% of the section blob.
280        modified_index = len(blob) / 50
281        if corrupt_all:
282            blob_list = [('%c' % ((ord(x) + delta) % 0x100)) for x in blob]
283        else:
284            blob_list = list(blob)
285            blob_list[modified_index] = ('%c' %
286                    ((ord(blob[modified_index]) + delta) % 0x100))
287        self.image = self.fum.put_section(self.image,
288                                          subsection_name, ''.join(blob_list))
289
290        return subsection_name
291
292    def corrupt_section(self, section, corrupt_all=False):
293        """Corrupt a section signature of the image"""
294
295        return self._modify_section(section, self.DELTA, body_or_sig=False,
296                                    corrupt_all=corrupt_all)
297
298    def corrupt_section_body(self, section, corrupt_all=False):
299        """Corrupt a section body of the image"""
300
301        return self._modify_section(section, self.DELTA, body_or_sig=True,
302                                    corrupt_all=corrupt_all)
303
304    def restore_section(self, section, restore_all=False):
305        """Restore a previously corrupted section signature of the image."""
306
307        return self._modify_section(section, -self.DELTA, body_or_sig=False,
308                                    corrupt_all=restore_all)
309
310    def restore_section_body(self, section, restore_all=False):
311        """Restore a previously corrupted section body of the image."""
312
313        return self._modify_section(section, -self.DELTA, body_or_sig=True,
314                                    corrupt_all=restore_all)
315
316    def corrupt_firmware(self, section, corrupt_all=False):
317        """Corrupt a section signature in the FLASHROM!!!"""
318
319        subsection_name = self.corrupt_section(section, corrupt_all=corrupt_all)
320        self.fum.write_partial(self.image, (subsection_name, ))
321
322    def corrupt_firmware_body(self, section, corrupt_all=False):
323        """Corrupt a section body in the FLASHROM!!!"""
324
325        subsection_name = self.corrupt_section_body(section,
326                                                    corrupt_all=corrupt_all)
327        self.fum.write_partial(self.image, (subsection_name, ))
328
329    def restore_firmware(self, section, restore_all=False):
330        """Restore the previously corrupted section sig in the FLASHROM!!!"""
331
332        subsection_name = self.restore_section(section, restore_all=restore_all)
333        self.fum.write_partial(self.image, (subsection_name, ))
334
335    def restore_firmware_body(self, section, restore_all=False):
336        """Restore the previously corrupted section body in the FLASHROM!!!"""
337
338        subsection_name = self.restore_section_body(section,
339                                                    restore_all=False)
340        self.fum.write_partial(self.image, (subsection_name, ))
341
342    def firmware_sections_equal(self):
343        """Check if firmware sections A and B are equal.
344
345        This function presumes that the entire BIOS image integrity has been
346        verified, so different signature sections mean different images and
347        vice versa.
348        """
349        sig_a = self.fum.get_section(self.image,
350                                      self.fv_sections['a'].get_sig_name())
351        sig_b = self.fum.get_section(self.image,
352                                      self.fv_sections['b'].get_sig_name())
353        return sig_a == sig_b
354
355    def copy_from_to(self, src, dst):
356        """Copy one firmware image section to another.
357
358        This function copies both signature and body of one firmware section
359        into another. After this function runs both sections are identical.
360        """
361        src_sect = self.fv_sections[src]
362        dst_sect = self.fv_sections[dst]
363        self.image = self.fum.put_section(
364            self.image,
365            dst_sect.get_body_name(),
366            self.fum.get_section(self.image, src_sect.get_body_name()))
367        self.image = self.fum.put_section(
368            self.image,
369            dst_sect.get_sig_name(),
370            self.fum.get_section(self.image, src_sect.get_sig_name()))
371
372    def write_whole(self):
373        """Write the whole image into the flashrom."""
374
375        if not self.image:
376            raise FlashromHandlerError(
377                'Attempt at using an uninitialized object')
378        self.fum.write_whole(self.image)
379
380    def dump_whole(self, filename):
381        """Write the whole image into a file."""
382
383        if not self.image:
384            raise FlashromHandlerError(
385                'Attempt at using an uninitialized object')
386        open(filename, 'w').write(self.image)
387
388    def dump_partial(self, subsection_name, filename):
389        """Write the subsection part into a file."""
390
391        if not self.image:
392            raise FlashromHandlerError(
393                'Attempt at using an uninitialized object')
394        blob = self.fum.get_section(self.image, subsection_name)
395        open(filename, 'w').write(blob)
396
397    def get_gbb_flags(self):
398        """Retrieve the GBB flags"""
399        gbb_header_format = '<12sL'
400        gbb_section = self.fum.get_section(self.image, 'FV_GBB')
401        try:
402            _, gbb_flags = struct.unpack_from(gbb_header_format, gbb_section)
403        except struct.error, e:
404            raise FlashromHandlerError(e)
405        return gbb_flags
406
407    def set_gbb_flags(self, flags, write_through=False):
408        """Retrieve the GBB flags"""
409        gbb_header_format = '<L'
410        section_name = 'FV_GBB'
411        gbb_section = self.fum.get_section(self.image, section_name)
412        try:
413            formatted_flags = struct.pack(gbb_header_format, flags)
414        except struct.error, e:
415            raise FlashromHandlerError(e)
416        gbb_section = gbb_section[:12] + formatted_flags + gbb_section[16:]
417        self.image = self.fum.put_section(self.image, section_name, gbb_section)
418
419        if write_through:
420            self.dump_partial(section_name,
421                              self.os_if.state_dir_file(section_name))
422            self.fum.write_partial(self.image, (section_name, ))
423
424    def enable_write_protect(self):
425        """Enable write protect of the flash chip"""
426        self.fum.enable_write_protect()
427
428    def disable_write_protect(self):
429        """Disable write protect of the flash chip"""
430        self.fum.disable_write_protect()
431
432    def get_section_sig_sha(self, section):
433        """Retrieve SHA1 hash of a firmware vblock section"""
434        return self.fv_sections[section].get_sig_sha()
435
436    def get_section_sha(self, section):
437        """Retrieve SHA1 hash of a firmware body section"""
438        return self.fv_sections[section].get_sha()
439
440    def get_section_version(self, section):
441        """Retrieve version number of a firmware section"""
442        return self.fv_sections[section].get_version()
443
444    def get_section_flags(self, section):
445        """Retrieve preamble flags of a firmware section"""
446        return self.fv_sections[section].get_flags()
447
448    def get_section_datakey_version(self, section):
449        """Retrieve data key version number of a firmware section"""
450        return self.fv_sections[section].get_datakey_version()
451
452    def get_section_kernel_subkey_version(self, section):
453        """Retrieve kernel subkey version number of a firmware section"""
454        return self.fv_sections[section].get_kernel_subkey_version()
455
456    def get_section_body(self, section):
457        """Retrieve body of a firmware section"""
458        subsection_name = self.fv_sections[section].get_body_name()
459        blob = self.fum.get_section(self.image, subsection_name)
460        return blob
461
462    def get_section_sig(self, section):
463        """Retrieve vblock of a firmware section"""
464        subsection_name = self.fv_sections[section].get_sig_name()
465        blob = self.fum.get_section(self.image, subsection_name)
466        return blob
467
468    def set_section_body(self, section, blob, write_through=False):
469        """Put the supplied blob to the body of the firmware section"""
470        subsection_name = self.fv_sections[section].get_body_name()
471        self.image = self.fum.put_section(self.image, subsection_name, blob)
472
473        if write_through:
474            self.dump_partial(subsection_name,
475                              self.os_if.state_dir_file(subsection_name))
476            self.fum.write_partial(self.image, (subsection_name, ))
477
478    def set_section_sig(self, section, blob, write_through=False):
479        """Put the supplied blob to the vblock of the firmware section"""
480        subsection_name = self.fv_sections[section].get_sig_name()
481        self.image = self.fum.put_section(self.image, subsection_name, blob)
482
483        if write_through:
484            self.dump_partial(subsection_name,
485                              self.os_if.state_dir_file(subsection_name))
486            self.fum.write_partial(self.image, (subsection_name, ))
487
488    def set_section_version(self, section, version, flags,
489                            write_through=False):
490        """
491        Re-sign the firmware section using the supplied version number and
492        flag.
493        """
494        if (self.get_section_version(section) == version and
495            self.get_section_flags(section) == flags):
496            return  # No version or flag change, nothing to do.
497        if version < 0:
498            raise FlashromHandlerError(
499                'Attempt to set version %d on section %s' % (version, section))
500        fv_section = self.fv_sections[section]
501        sig_name = self.os_if.state_dir_file(fv_section.get_sig_name())
502        sig_size = os.path.getsize(sig_name)
503
504        # Construct the command line
505        args = ['--vblock %s' % sig_name]
506        args.append('--keyblock %s' % os.path.join(
507                self.dev_key_path, self.FW_KEYBLOCK_FILE_NAME))
508        args.append('--fv %s' % self.os_if.state_dir_file(
509                fv_section.get_body_name()))
510        args.append('--version %d' % version)
511        args.append('--kernelkey %s' % os.path.join(
512                self.dev_key_path, self.KERNEL_SUBKEY_FILE_NAME))
513        args.append('--signprivate %s' % os.path.join(
514                self.dev_key_path, self.FW_PRIV_DATA_KEY_FILE_NAME))
515        args.append('--flags %d' % flags)
516        cmd = 'vbutil_firmware %s' % ' '.join(args)
517        self.os_if.run_shell_command(cmd)
518
519        #  Pad the new signature.
520        new_sig = open(sig_name, 'a')
521        pad = ('%c' % 0) * (sig_size - os.path.getsize(sig_name))
522        new_sig.write(pad)
523        new_sig.close()
524
525        # Inject the new signature block into the image
526        new_sig = open(sig_name, 'r').read()
527        self.image = self.fum.put_section(
528            self.image, fv_section.get_sig_name(), new_sig)
529        if write_through:
530            self.fum.write_partial(self.image, (fv_section.get_sig_name(), ))
531