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