• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/python2
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"""A module to support automated testing of ChromeOS firmware.
6
7Utilizes services provided by saft_flashrom_util.py read/write the
8flashrom chip and to parse the flash rom image.
9
10See docstring for FlashromHandler class below.
11"""
12
13import hashlib
14import os
15import struct
16import tempfile
17
18from autotest_lib.client.common_lib import error
19from autotest_lib.client.common_lib.cros import chip_utils
20from autotest_lib.client.cros.faft.utils import saft_flashrom_util
21
22
23class FvSection(object):
24    """An object to hold information about a firmware section.
25
26    This includes file names for the signature header and the body, and the
27    version number.
28    """
29
30    def __init__(self, sig_name, body_name, fwid_name=None):
31        """
32        @param sig_name: name of signature section in fmap
33        @param body_name: name of body section in fmap
34        @param fwid_name: name of fwid section in fmap
35        @type sig_name: str | None
36        @type body_name: str | None
37        @type fwid_name: str | None
38        """
39        self._sig_name = sig_name
40        self._body_name = body_name
41        self._fwid_name = fwid_name
42        self._version = -1  # Is not set on construction.
43        self._flags = 0  # Is not set on construction.
44        self._sha = None  # Is not set on construction.
45        self._sig_sha = None  # Is not set on construction.
46        self._datakey_version = -1  # Is not set on construction.
47        self._kernel_subkey_version = -1  # Is not set on construction.
48
49    def names(self):
50        """Return the desired file names for the signature, body, and fwid."""
51        return (self._sig_name, self._body_name, self._fwid_name)
52
53    def get_sig_name(self):
54        return self._sig_name
55
56    def get_body_name(self):
57        return self._body_name
58
59    def get_fwid_name(self):
60        return self._fwid_name
61
62    def get_version(self):
63        return self._version
64
65    def get_flags(self):
66        return self._flags
67
68    def get_sha(self):
69        return self._sha
70
71    def get_sig_sha(self):
72        return self._sig_sha
73
74    def get_datakey_version(self):
75        return self._datakey_version
76
77    def get_kernel_subkey_version(self):
78        return self._kernel_subkey_version
79
80    def set_version(self, version):
81        self._version = version
82
83    def set_flags(self, flags):
84        self._flags = flags
85
86    def set_sha(self, sha):
87        self._sha = sha
88
89    def set_sig_sha(self, sha):
90        self._sig_sha = sha
91
92    def set_datakey_version(self, version):
93        self._datakey_version = version
94
95    def set_kernel_subkey_version(self, version):
96        self._kernel_subkey_version = version
97
98
99class FlashromHandlerError(Exception):
100    """An object to represent Flashrom errors"""
101    pass
102
103
104class FlashromHandler(object):
105    """An object to provide logical services for automated flashrom testing."""
106
107    DELTA = 1  # value to add to a byte to corrupt a section contents
108
109    # File in the state directory to store public root key.
110    PUB_KEY_FILE_NAME = 'root.pubkey'
111    FW_KEYBLOCK_FILE_NAME = 'firmware.keyblock'
112    FW_PRIV_DATA_KEY_FILE_NAME = 'firmware_data_key.vbprivk'
113    KERNEL_SUBKEY_FILE_NAME = 'kernel_subkey.vbpubk'
114    EC_EFS_KEY_FILE_NAME = 'key_ec_efs.vbprik2'
115    FWID_MOD_DELIMITER = '~'
116
117    def __init__(
118            self,
119            os_if,
120            pub_key_file=None,
121            dev_key_path='./',
122            target='bios',
123            subdir=None
124    ):
125        """The flashrom handler is not fully initialized upon creation
126
127        @param os_if: an object providing interface to OS services
128        @param pub_key_file: the name of the file contaning a public key to
129                             use for verifying both existing and new firmware.
130        @param dev_key_path: path to directory containing *.vpubk and *.vbprivk
131                             files, for use in signing
132        @param target: flashrom target ('bios' or 'ec')
133        @param subdir: name of subdirectory of state dir, to use for sections
134                    Default: same as target, resulting in
135                    '/usr/local/tmp/faft/bios'
136        @type os_if: client.cros.faft.utils.os_interface.OSInterface
137        @type pub_key_file: str | None
138        @type dev_key_path: str
139        @type target: str
140        """
141        self.fum = None
142        self.image = ''
143        self.os_if = os_if
144        self.initialized = False
145        self._available = None
146        self._unavailable_err = None
147
148        if subdir is None:
149            subdir = target
150        self.subdir = subdir
151
152        self.pub_key_file = pub_key_file
153        self.dev_key_path = dev_key_path
154
155        self.target = target
156        if self.target == 'bios':
157            self.fum = saft_flashrom_util.flashrom_util(
158                    self.os_if, target_is_ec=False)
159            self.fv_sections = {
160                    'ro': FvSection(None, None, 'RO_FRID'),
161                    'a': FvSection('VBOOTA', 'FVMAIN', 'RW_FWID_A'),
162                    'b': FvSection('VBOOTB', 'FVMAINB', 'RW_FWID_B'),
163                    'rec': FvSection(None, 'RECOVERY_MRC_CACHE'),
164                    'ec_a': FvSection(None, 'ECMAINA'),
165                    'ec_b': FvSection(None, 'ECMAINB'),
166                    'rw_legacy': FvSection(None, 'RW_LEGACY'),
167            }
168        elif self.target == 'ec':
169            self.fum = saft_flashrom_util.flashrom_util(
170                    self.os_if, target_is_ec=True)
171            self.fv_sections = {
172                    'ro': FvSection(None, None, 'RO_FRID'),
173                    'rw': FvSection(None, 'EC_RW', 'RW_FWID'),
174                    'rw_b': FvSection(None, 'EC_RW_B'),
175            }
176        else:
177            raise FlashromHandlerError("Invalid target.")
178
179    def is_available(self):
180        """Check if the programmer is available, by specifying no commands.
181
182        @rtype: bool
183        """
184        if self._available is None:
185            # Cache the status to avoid trying flashrom every time.
186            try:
187                self.fum.check_target()
188                self._available = True
189            except error.CmdError as e:
190                # First line: "Command <flashrom -p host> failed, rc=2"
191                self._unavailable_err = str(e).split('\n', 1)[0]
192                self._available = False
193        return self._available
194
195    def section_file(self, *paths):
196        """
197        Return a full path for the given basename, in this handler's subdir.
198        Example: subdir 'bios' -> '/usr/local/tmp/faft/bios/FV_GBB'
199
200        @param paths: variable number of path pieces, same as in os.path.join
201        @return: an absolute path from this handler's subdir and the pieces.
202        """
203        if any(os.path.isabs(x) for x in paths):
204            raise FlashromHandlerError(
205                    "Absolute paths are not allowed in section_file()")
206
207        return os.path.join(self.os_if.state_dir, self.subdir, *paths)
208
209    def init(self, image_file=None, allow_fallback=False):
210        """Initialize the object, by reading the image.
211
212        This is separate from new_image, to isolate the implementation detail of
213        self.image being non-empty.
214
215        @param image_file: the path of the image file to read.
216                If None or empty string, read the flash device instead.
217        @param allow_fallback: if True, fall back to reading the flash device
218                if the image file doesn't exist.
219        @type image_file: str
220        @type allow_fallback: bool
221
222        @raise FlashromHandlerError: if no target flash device was usable.
223        """
224        # Raise an exception early if there's no usable flash.
225        if not self.is_available():
226            # Can't tell for sure whether it's broken or simply nonexistent.
227            raise FlashromHandlerError(
228                    "Could not detect a usable %s flash device: %s."
229                    % (self.target, self._unavailable_err))
230
231        if image_file and allow_fallback and not os.path.isfile(image_file):
232            self.os_if.log(
233                    "Using %s flash contents instead of missing image: %s"
234                    % (self.target.upper(), image_file))
235            image_file = None
236
237        self.new_image(image_file)
238        self.initialized = True
239
240    def deinit(self):
241        """Clear the in-memory image data, and mark self uninitialized."""
242        self.image = ''
243        self.os_if.remove_dir(self.section_file())
244        self.initialized = False
245
246    def dump_flash(self, target_filename):
247        """Copy the flash device's data into a file, but don't parse it.
248
249        @param target_filename: the file to create
250        """
251        self.fum.dump_flash(target_filename)
252
253    def new_image(self, image_file=None):
254        """Parse the full flashrom image and store sections into files.
255
256        @param image_file: the name of the file containing a full ChromeOS
257                       flashrom image. If not passed in or empty, the actual
258                       flash device is read and its contents are saved into a
259                       temporary file which is used instead.
260        @type image_file: str | None
261
262        The input file is parsed and the sections of importance (as defined in
263        self.fv_sections) are saved in separate files in the state directory
264        as defined in the os_if object.
265        """
266
267        if image_file:
268            with open(image_file, 'rb') as image_f:
269                self.image = image_f.read()
270            self.fum.set_firmware_layout(image_file)
271        else:
272            self.image = self.fum.read_whole()
273
274        self.os_if.create_dir(self.section_file())
275
276        for section in self.fv_sections.itervalues():
277            for subsection_name in section.names():
278                if not subsection_name:
279                    continue
280                blob = self.fum.get_section(self.image, subsection_name)
281                if blob:
282                    blob_filename = self.section_file(subsection_name)
283                    with open(blob_filename, 'wb') as blob_f:
284                        blob_f.write(blob)
285
286            blob = self.fum.get_section(self.image, section.get_body_name())
287            if blob:
288                s = hashlib.sha1()
289                s.update(blob)
290                section.set_sha(s.hexdigest())
291
292            # If there is no "sig" subsection, skip reading version and flags.
293            if not section.get_sig_name():
294                continue
295
296            # Now determine this section's version number.
297            vb_section = self.fum.get_section(self.image,
298                                              section.get_sig_name())
299
300            section.set_version(self.os_if.retrieve_body_version(vb_section))
301            section.set_flags(self.os_if.retrieve_preamble_flags(vb_section))
302            section.set_datakey_version(
303                    self.os_if.retrieve_datakey_version(vb_section))
304            section.set_kernel_subkey_version(
305                    self.os_if.retrieve_kernel_subkey_version(vb_section))
306
307            s = hashlib.sha1()
308            s.update(self.fum.get_section(self.image, section.get_sig_name()))
309            section.set_sig_sha(s.hexdigest())
310
311        if not self.pub_key_file:
312            self._retrieve_pub_key()
313
314    def _retrieve_pub_key(self):
315        """Retrieve root public key from the firmware GBB section."""
316
317        gbb_header_format = '<4s20s2I'
318        pubk_header_format = '<2Q'
319
320        gbb_section = self.fum.get_section(self.image, 'FV_GBB')
321
322        # do some sanity checks
323        try:
324            sig, _, rootk_offs, rootk_size = struct.unpack_from(
325                    gbb_header_format, gbb_section)
326        except struct.error, e:
327            raise FlashromHandlerError(e)
328
329        if sig != '$GBB' or (rootk_offs + rootk_size) > len(gbb_section):
330            raise FlashromHandlerError('Bad gbb header')
331
332        key_body_offset, key_body_size = struct.unpack_from(
333                pubk_header_format, gbb_section, rootk_offs)
334
335        # Generally speaking the offset field can be anything, but in case of
336        # GBB section the key is stored as a standalone entity, so the offset
337        # of the key body is expected to be equal to the key header size of
338        # 0x20.
339        # Should this convention change, the check below would fail, which
340        # would be a good prompt for revisiting this test's behavior and
341        # algorithms.
342        if key_body_offset != 0x20 or key_body_size > rootk_size:
343            raise FlashromHandlerError('Bad public key format')
344
345        # All checks passed, let's store the key in a file.
346        self.pub_key_file = self.os_if.state_dir_file(self.PUB_KEY_FILE_NAME)
347        with open(self.pub_key_file, 'w') as key_f:
348            key = gbb_section[rootk_offs:rootk_offs + key_body_offset +
349                              key_body_size]
350            key_f.write(key)
351
352    def verify_image(self):
353        """Confirm the image's validity.
354
355        Using the file supplied to init() as the public key container verify
356        the two sections' (FirmwareA and FirmwareB) integrity. The contents of
357        the sections is taken from the files created by new_image()
358
359        In case there is an integrity error raises FlashromHandlerError
360        exception with the appropriate error message text.
361        """
362
363        for section in self.fv_sections.itervalues():
364            if section.get_sig_name():
365                cmd = 'vbutil_firmware --verify %s --signpubkey %s  --fv %s' % (
366                        self.section_file(section.get_sig_name()),
367                        self.pub_key_file,
368                        self.section_file(section.get_body_name()))
369                self.os_if.run_shell_command(cmd)
370
371    def _modify_section(self,
372                        section,
373                        delta,
374                        body_or_sig=False,
375                        corrupt_all=False):
376        """Modify a firmware section inside the image, either body or signature.
377
378        If corrupt_all is set, the passed in delta is added to all bytes in the
379        section. Otherwise, the delta is added to the value located at 2% offset
380        into the section blob, either body or signature.
381
382        Calling this function again for the same section the complimentary
383        delta value would restore the section contents.
384        """
385
386        if not self.image:
387            raise FlashromHandlerError(
388                    'Attempt at using an uninitialized object')
389        if section not in self.fv_sections:
390            raise FlashromHandlerError('Unknown FW section %s' % section)
391
392        # Get the appropriate section of the image.
393        if body_or_sig:
394            subsection_name = self.fv_sections[section].get_body_name()
395        else:
396            subsection_name = self.fv_sections[section].get_sig_name()
397        blob = self.fum.get_section(self.image, subsection_name)
398
399        # Modify the byte in it within 2% of the section blob.
400        modified_index = len(blob) / 50
401        if corrupt_all:
402            blob_list = [('%c' % ((ord(x) + delta) % 0x100)) for x in blob]
403        else:
404            blob_list = list(blob)
405            blob_list[modified_index] = (
406                    '%c' % ((ord(blob[modified_index]) + delta) % 0x100))
407        self.image = self.fum.put_section(self.image, subsection_name,
408                                          ''.join(blob_list))
409
410        return subsection_name
411
412    def corrupt_section(self, section, corrupt_all=False):
413        """Corrupt a section signature of the image"""
414
415        return self._modify_section(
416                section,
417                self.DELTA,
418                body_or_sig=False,
419                corrupt_all=corrupt_all)
420
421    def corrupt_section_body(self, section, corrupt_all=False):
422        """Corrupt a section body of the image"""
423
424        return self._modify_section(
425                section, self.DELTA, body_or_sig=True, corrupt_all=corrupt_all)
426
427    def restore_section(self, section, restore_all=False):
428        """Restore a previously corrupted section signature of the image."""
429
430        return self._modify_section(
431                section,
432                -self.DELTA,
433                body_or_sig=False,
434                corrupt_all=restore_all)
435
436    def restore_section_body(self, section, restore_all=False):
437        """Restore a previously corrupted section body of the image."""
438
439        return self._modify_section(
440                section,
441                -self.DELTA,
442                body_or_sig=True,
443                corrupt_all=restore_all)
444
445    def corrupt_firmware(self, section, corrupt_all=False):
446        """Corrupt a section signature in the FLASHROM!!!"""
447
448        subsection_name = self.corrupt_section(
449                section, corrupt_all=corrupt_all)
450        self.fum.write_partial(self.image, (subsection_name, ))
451
452    def corrupt_firmware_body(self, section, corrupt_all=False):
453        """Corrupt a section body in the FLASHROM!!!"""
454
455        subsection_name = self.corrupt_section_body(
456                section, corrupt_all=corrupt_all)
457        self.fum.write_partial(self.image, (subsection_name, ))
458
459    def restore_firmware(self, section, restore_all=False):
460        """Restore the previously corrupted section sig in the FLASHROM!!!"""
461
462        subsection_name = self.restore_section(
463                section, restore_all=restore_all)
464        self.fum.write_partial(self.image, (subsection_name, ))
465
466    def restore_firmware_body(self, section, restore_all=False):
467        """Restore the previously corrupted section body in the FLASHROM!!!"""
468
469        subsection_name = self.restore_section_body(section, restore_all=False)
470        self.fum.write_partial(self.image, (subsection_name, ))
471
472    def firmware_sections_equal(self):
473        """Check if firmware sections A and B are equal.
474
475        This function presumes that the entire BIOS image integrity has been
476        verified, so different signature sections mean different images and
477        vice versa.
478        """
479        sig_a = self.fum.get_section(self.image,
480                                     self.fv_sections['a'].get_sig_name())
481        sig_b = self.fum.get_section(self.image,
482                                     self.fv_sections['b'].get_sig_name())
483        return sig_a == sig_b
484
485    def copy_from_to(self, src, dst):
486        """Copy one firmware image section to another.
487
488        This function copies both signature and body of one firmware section
489        into another. After this function runs both sections are identical.
490        """
491        src_sect = self.fv_sections[src]
492        dst_sect = self.fv_sections[dst]
493        self.image = self.fum.put_section(
494                self.image, dst_sect.get_body_name(),
495                self.fum.get_section(self.image, src_sect.get_body_name()))
496        # If there is no "sig" subsection, skip copying signature.
497        if src_sect.get_sig_name() and dst_sect.get_sig_name():
498            self.image = self.fum.put_section(
499                    self.image, dst_sect.get_sig_name(),
500                    self.fum.get_section(self.image, src_sect.get_sig_name()))
501        self.write_whole()
502
503    def write_whole(self):
504        """Write the whole image into the flashrom."""
505
506        if not self.image:
507            raise FlashromHandlerError(
508                    'Attempt at using an uninitialized object')
509        self.fum.write_whole(self.image)
510
511    def write_partial(self, subsection_name, blob=None, write_through=True):
512        """Write the subsection part into the flashrom.
513
514        One can pass a blob to update the data of the subsection before write
515        it into the flashrom.
516        """
517
518        if not self.image:
519            raise FlashromHandlerError(
520                    'Attempt at using an uninitialized object')
521
522        if blob is not None:
523            self.image = self.fum.put_section(self.image, subsection_name,
524                                              blob)
525
526        if write_through:
527            self.dump_partial(
528                    subsection_name, self.section_file(subsection_name))
529            self.fum.write_partial(self.image, (subsection_name, ))
530
531    def dump_whole(self, filename):
532        """Write the whole image into a file."""
533
534        if not self.image:
535            raise FlashromHandlerError(
536                    'Attempt at using an uninitialized object')
537        open(filename, 'w').write(self.image)
538
539    def dump_partial(self, subsection_name, filename):
540        """Write the subsection part into a file."""
541
542        if not self.image:
543            raise FlashromHandlerError(
544                    'Attempt at using an uninitialized object')
545        blob = self.fum.get_section(self.image, subsection_name)
546        open(filename, 'w').write(blob)
547
548    def dump_section_body(self, section, filename):
549        """Write the body of a firmware section into a file"""
550        subsection_name = self.fv_sections[section].get_body_name()
551        self.dump_partial(subsection_name, filename)
552
553    def get_section_hash(self, section):
554        """Retrieve the hash of the body of a firmware section"""
555        ecrw = chip_utils.ecrw()
556
557        # add a dot to avoid set_from_file breaking if tmpname has an underscore
558        with tempfile.NamedTemporaryFile(prefix=ecrw.chip_name + '.') as f:
559            self.dump_section_body(section, f.name)
560            ecrw.set_from_file(f.name)
561            result = ecrw.compute_hash_bytes()
562        return result
563
564    def get_gbb_flags(self):
565        """Retrieve the GBB flags"""
566        gbb_header_format = '<12sL'
567        gbb_section = self.fum.get_section(self.image, 'FV_GBB')
568        try:
569            _, gbb_flags = struct.unpack_from(gbb_header_format, gbb_section)
570        except struct.error, e:
571            raise FlashromHandlerError(e)
572        return gbb_flags
573
574    def set_gbb_flags(self, flags, write_through=False):
575        """Retrieve the GBB flags"""
576        gbb_header_format = '<L'
577        section_name = 'FV_GBB'
578        gbb_section = self.fum.get_section(self.image, section_name)
579        try:
580            formatted_flags = struct.pack(gbb_header_format, flags)
581        except struct.error, e:
582            raise FlashromHandlerError(e)
583        gbb_section = gbb_section[:12] + formatted_flags + gbb_section[16:]
584        self.write_partial(section_name, gbb_section, write_through)
585
586    def enable_write_protect(self):
587        """Enable write protect of the flash chip"""
588        self.fum.enable_write_protect()
589
590    def disable_write_protect(self):
591        """Disable write protect of the flash chip"""
592        self.fum.disable_write_protect()
593
594    def set_write_protect_region(self, region, enabled=None):
595        """
596        Set write protection region by name, using current image's layout.
597
598        The name should match those seen in `futility dump_fmap <image>`, and
599        is not checked against self.firmware_layout, due to different naming.
600
601        @param region: Region to set (usually WP_RO)
602        @param enabled: If True, run --wp-enable; if False, run --wp-disable.
603                        If None (default), don't specify either one.
604        """
605        if region is None:
606            raise FlashromHandlerError("Region must not be None")
607        image_file = self.os_if.create_temp_file('wp_')
608        self.os_if.write_file(image_file, self.image)
609
610        self.fum.set_write_protect_region(image_file, region, enabled)
611        self.os_if.remove_file(image_file)
612
613    def set_write_protect_range(self, start, length, enabled=None):
614        """
615        Set write protection range by offsets, using current image's layout.
616
617        @param start: offset (bytes) from start of flash to start of range
618        @param length: offset (bytes) from start of range to end of range
619        @param enabled: If True, run --wp-enable; if False, run --wp-disable.
620                        If None (default), don't specify either one.
621        """
622        self.fum.set_write_protect_range(start, length, enabled)
623
624    def get_write_protect_status(self):
625        """Get a dict describing the status of the write protection
626
627        @return: {'enabled': True/False, 'start': '0x0', 'length': '0x0', ...}
628        @rtype: dict
629        """
630        return self.fum.get_write_protect_status()
631
632    def get_section_sig_sha(self, section):
633        """Retrieve SHA1 hash of a firmware vblock section"""
634        return self.fv_sections[section].get_sig_sha()
635
636    def get_section_sha(self, section):
637        """Retrieve SHA1 hash of a firmware body section"""
638        return self.fv_sections[section].get_sha()
639
640    def get_section_version(self, section):
641        """Retrieve version number of a firmware section"""
642        return self.fv_sections[section].get_version()
643
644    def get_section_flags(self, section):
645        """Retrieve preamble flags of a firmware section"""
646        return self.fv_sections[section].get_flags()
647
648    def get_section_datakey_version(self, section):
649        """Retrieve data key version number of a firmware section"""
650        return self.fv_sections[section].get_datakey_version()
651
652    def get_section_kernel_subkey_version(self, section):
653        """Retrieve kernel subkey version number of a firmware section"""
654        return self.fv_sections[section].get_kernel_subkey_version()
655
656    def get_section_body(self, section):
657        """Retrieve body of a firmware section"""
658        subsection_name = self.fv_sections[section].get_body_name()
659        blob = self.fum.get_section(self.image, subsection_name)
660        return blob
661
662    def has_section_body(self, section):
663        """Return True if the section body is in the image"""
664        return bool(self.get_section_body(section))
665
666    def get_section_sig(self, section):
667        """Retrieve vblock of a firmware section"""
668        subsection_name = self.fv_sections[section].get_sig_name()
669        blob = self.fum.get_section(self.image, subsection_name)
670        return blob
671
672    def get_section_fwid(self, section, strip_null=True):
673        """
674        Retrieve fwid blob of a firmware section.
675
676        @param section: Name of the section whose fwid to return.
677        @param strip_null: If True, remove \0 from the end of the blob.
678        @return: fwid of the section
679
680        @type section: str
681        @type strip_null: bool
682        @rtype: str | None
683
684        """
685        subsection_name = self.fv_sections[section].get_fwid_name()
686        if not subsection_name:
687            return None
688        blob = self.fum.get_section(self.image, subsection_name)
689        if strip_null:
690            blob = blob.rstrip('\0')
691        return blob
692
693    def set_section_body(self, section, blob, write_through=False):
694        """Put the supplied blob to the body of the firmware section"""
695        subsection_name = self.fv_sections[section].get_body_name()
696        self.write_partial(subsection_name, blob, write_through)
697
698    def set_section_sig(self, section, blob, write_through=False):
699        """Put the supplied blob to the vblock of the firmware section"""
700        subsection_name = self.fv_sections[section].get_sig_name()
701        self.write_partial(subsection_name, blob, write_through)
702
703    def set_section_fwid(self, section, blob, write_through=False):
704        """Put the supplied blob to the fwid of the firmware section"""
705        subsection_name = self.fv_sections[section].get_fwid_name()
706        self.write_partial(subsection_name, blob, write_through)
707
708    def resign_ec_rwsig(self):
709        """Resign the EC image using rwsig."""
710        key_ec_efs = os.path.join(self.dev_key_path, self.EC_EFS_KEY_FILE_NAME)
711        # Dump whole EC image to a file and execute the sign command.
712        with tempfile.NamedTemporaryFile() as f:
713            self.dump_whole(f.name)
714            self.os_if.run_shell_command(
715                    'futility sign --type rwsig --prikey %s %s' % (key_ec_efs,
716                                                                   f.name))
717            self.new_image(f.name)
718
719    def set_section_version(self, section, version, flags,
720                            write_through=False):
721        """
722        Re-sign the firmware section using the supplied version number and
723        flag.
724        """
725        if (self.get_section_version(section) == version
726                    and self.get_section_flags(section) == flags):
727            return  # No version or flag change, nothing to do.
728        if version < 0:
729            raise FlashromHandlerError(
730                    'Attempt to set version %d on section %s' % (version,
731                                                                 section))
732        fv_section = self.fv_sections[section]
733        sig_name = self.section_file(fv_section.get_sig_name())
734        sig_size = os.path.getsize(sig_name)
735
736        # Construct the command line
737        args = ['--vblock %s' % sig_name]
738        args.append('--keyblock %s' % os.path.join(self.dev_key_path,
739                                                   self.FW_KEYBLOCK_FILE_NAME))
740        args.append('--fv %s' % self.section_file(fv_section.get_body_name()))
741        args.append('--version %d' % version)
742        args.append('--kernelkey %s' % os.path.join(
743                self.dev_key_path, self.KERNEL_SUBKEY_FILE_NAME))
744        args.append('--signprivate %s' % os.path.join(
745                self.dev_key_path, self.FW_PRIV_DATA_KEY_FILE_NAME))
746        args.append('--flags %d' % flags)
747        cmd = 'vbutil_firmware %s' % ' '.join(args)
748        self.os_if.run_shell_command(cmd)
749
750        #  Pad the new signature.
751        with open(sig_name, 'a') as sig_f:
752            f_size = os.fstat(sig_f.fileno()).st_size
753            pad = '\0' * (sig_size - f_size)
754            sig_f.write(pad)
755
756        # Inject the new signature block into the image
757        with open(sig_name, 'r') as sig_f:
758            new_sig = sig_f.read()
759        self.write_partial(fv_section.get_sig_name(), new_sig, write_through)
760
761    def _modify_section_fwid(self, section):
762        """Modify a section's fwid on the handler, adding a tilde and the
763        section name (in caps) to the end: ~RO, ~RW, ~A, ~B.
764
765        @param section: the single section to act on
766        @return: the new fwid
767
768        @type section: str
769        @rtype: str
770        """
771
772        fwid = self.get_section_fwid(section, strip_null=False)
773
774        if fwid is None:
775            return None
776
777        fwid_size = len(fwid)
778
779        if not fwid:
780            raise FlashromHandlerError(
781                    "FWID (%s, %s) is empty: %s" %
782                    (self.target.upper(), section.upper(), repr(fwid)))
783
784        fwid = fwid.rstrip('\0')
785        suffix = self.FWID_MOD_DELIMITER + section.upper()
786
787        if suffix in fwid:
788            raise FlashromHandlerError(
789                    "FWID (%s, %s) is already modified: %s" %
790                    (self.target.upper(), section.upper(), repr(fwid)))
791
792        # Append a suffix, after possibly chopping off characters to make room.
793        if len(fwid) + len(suffix) > fwid_size:
794            fwid = fwid[:fwid_size - len(suffix)]
795        fwid += suffix
796
797        padded_fwid = fwid.ljust(fwid_size, '\0')
798        self.set_section_fwid(section, padded_fwid)
799        return fwid
800
801    def _strip_section_fwid(self, section, write_through=True):
802        """Modify a section's fwid on the handler, stripping any suffix added
803        by _modify_section_fwid: ~RO, ~RW, ~A, ~B.
804
805        @param section: the single section to act on
806        @param write_through: if True (default), write to flash immediately
807        @return: the suffix that was stripped
808
809        @type section: str
810        @type write_through: bool
811        @rtype: str | None
812        """
813
814        fwid = self.get_section_fwid(section, strip_null=False)
815        if fwid is None:
816            return None
817
818        fwid_size = len(fwid)
819
820        if not fwid:
821            raise FlashromHandlerError(
822                    "FWID (%s, %s) is empty: %s" %
823                    (self.target.upper(), section.upper(), repr(fwid)))
824
825        fwid = fwid.rstrip('\0')
826        mod_indicator = self.FWID_MOD_DELIMITER + section.upper()
827
828        # Remove any suffix, and return the suffix if found.
829        if mod_indicator in fwid:
830            (stripped_fwid, remainder) = fwid.split(mod_indicator, 1)
831
832            padded_fwid = stripped_fwid.ljust(fwid_size, '\0')
833            self.set_section_fwid(section, padded_fwid, write_through)
834
835            return fwid
836        return None
837
838    def modify_fwids(self, sections):
839        """Modify the fwid in the in-memory image.
840
841        @param sections: section(s) to modify.
842        @return: fwids for the modified sections, as {section: fwid}
843
844        @type sections: tuple | list
845        @rtype: dict
846        """
847        fwids = {}
848        for section in sections:
849            fwids[section] = self._modify_section_fwid(section)
850
851        return fwids
852
853    def strip_modified_fwids(self):
854        """Strip any trailing suffixes (from modify_fwids) out of the FWIDs.
855
856        @return: a dict of any fwids that were adjusted, by section (ro, a, b)
857        @rtype: dict
858        """
859
860        suffixes = {}
861        for section in self.fv_sections:
862            suffix = self._strip_section_fwid(section)
863            if suffix is not None:
864                suffixes[section] = suffix
865
866        return suffixes
867