1# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4""" This module provides convenience routines to access Flash ROM (EEPROM)
5
6saft_flashrom_util is based on utility 'flashrom'.
7
8Original tool syntax:
9    (read ) flashrom -r <file>
10    (write) flashrom -l <layout_fn> [-i <image_name> ...] -w <file>
11
12The layout_fn is in format of
13    address_begin:address_end image_name
14    which defines a region between (address_begin, address_end) and can
15    be accessed by the name image_name.
16
17Currently the tool supports multiple partial write but not partial read.
18
19In the saft_flashrom_util, we provide read and partial write abilities.
20For more information, see help(saft_flashrom_util.flashrom_util).
21"""
22import re
23
24
25class TestError(Exception):
26    """Represents an internal error, such as invalid arguments."""
27    pass
28
29
30class LayoutScraper(object):
31    """Object of this class is used to retrieve layout from a BIOS file."""
32
33    DEFAULT_CHROMEOS_FMAP_CONVERSION = {
34            "BOOT_STUB": "FV_BSTUB",
35            "RO_FRID": "RO_FRID",
36            "GBB": "FV_GBB",
37            "RECOVERY": "FVDEV",
38            "VBLOCK_A": "VBOOTA",
39            "VBLOCK_B": "VBOOTB",
40            "FW_MAIN_A": "FVMAIN",
41            "FW_MAIN_B": "FVMAINB",
42            "RW_FWID_A": "RW_FWID_A",
43            "RW_FWID_B": "RW_FWID_B",
44            # Intel CSME FW Update sections
45            "ME_RW_A": "ME_RW_A",
46            "ME_RW_B": "ME_RW_B",
47            # Memory Training data cache for recovery boots
48            # Added on Nov 09, 2016
49            "RECOVERY_MRC_CACHE": "RECOVERY_MRC_CACHE",
50            # New sections in Depthcharge.
51            "EC_MAIN_A": "ECMAINA",
52            "EC_MAIN_B": "ECMAINB",
53            # EC firmware layout
54            "EC_RW": "EC_RW",
55            "EC_RW_B": "EC_RW_B",
56            "RW_FWID": "RW_FWID",
57            "RW_LEGACY": "RW_LEGACY",
58    }
59
60    def __init__(self, os_if):
61        self.image = None
62        self.os_if = os_if
63
64    def check_layout(self, layout, file_size):
65        """Verify the layout to be consistent.
66
67        The layout is consistent if there is no overlapping sections and the
68        section boundaries do not exceed the file size.
69
70        Inputs:
71          layout: a dictionary keyed by a string (the section name) with
72                  values being two integers tuples, the first and the last
73                  bites' offset in the file.
74          file_size: and integer, the size of the file the layout describes
75                     the sections in.
76
77        Raises:
78          TestError in case the layout is not consistent.
79        """
80
81        # Generate a list of section range tuples.
82        ost = sorted([layout[section] for section in layout])
83        base = -1
84        for section_base, section_end in ost:
85            if section_base <= base or section_end + 1 < section_base:
86                # Overlapped section is possible, like the fwid which is
87                # inside the main fw section.
88                self.os_if.log('overlapped section at 0x%x..0x%x' %
89                               (section_base, section_end))
90            base = section_end
91        if base > file_size:
92            raise TestError('Section end 0x%x exceeds file size %x' %
93                            (base, file_size))
94
95    def get_layout(self, file_name):
96        """Generate layout for a firmware file.
97
98        Internally, this uses the "dump_fmap" command, and converts
99        the output into a dictionary mapping region names to 2-tuples
100        of the start and last addresses.
101
102        Then verify the generated layout's consistency and return it to the
103        caller.
104        """
105        command = 'dump_fmap -p %s' % file_name
106        layout_data = {}  # keyed by the section name, elements - tuples of
107        # (<section start addr>, <section end addr>)
108
109        for line in self.os_if.run_shell_command_get_output(command):
110            region_name, offset, size = line.split()
111
112            try:
113                name = self.DEFAULT_CHROMEOS_FMAP_CONVERSION[region_name]
114            except KeyError:
115                continue  # This line does not contain an area of interest.
116
117            if name in layout_data:
118                raise TestError('%s duplicated in the layout' % name)
119
120            offset = int(offset)
121            size = int(size)
122            layout_data[name] = (offset, offset + size - 1)
123
124        self.check_layout(layout_data, self.os_if.get_file_size(file_name))
125        return layout_data
126
127
128# flashrom utility wrapper
129class flashrom_util(object):
130    """ a wrapper for "flashrom" utility.
131
132    You can read, write, or query flash ROM size with this utility.
133    Although you can do "partial-write", the tools always takes a
134    full ROM image as input parameter.
135
136    NOTE before accessing flash ROM, you may need to first "select"
137    your target - usually BIOS or EC. That part is not handled by
138    this utility. Please find other external script to do it.
139
140    To perform a read, you need to:
141     1. Prepare a flashrom_util object
142        ex: flashrom = flashrom_util.flashrom_util()
143     2. Perform read operation
144        ex: image = flashrom.read_whole()
145
146        When the contents of the flashrom is read off the target, it's map
147        gets created automatically (read from the flashrom image using
148        'dump_fmap'). If the user wants this object to operate on some other
149        file, they could either have the map for the file created explicitly by
150        invoking flashrom.set_firmware_layout(filename), or supply their own map
151        (which is a dictionary where keys are section names, and values are
152        tuples of integers, base address of the section and the last address
153        of the section).
154
155    By default this object operates on the map retrieved from the image and
156    stored locally, this map can be overwritten by an explicitly passed user
157    map.
158
159    To perform a (partial) write:
160
161     1. Prepare a buffer storing an image to be written into the flashrom.
162     2. Have the map generated automatically or prepare your own, for instance:
163        ex: layout_map_all = { 'all': (0, rom_size - 1) }
164        ex: layout_map = { 'ro': (0, 0xFFF), 'rw': (0x1000, rom_size-1) }
165     4. Perform write operation
166
167        ex using default map:
168          flashrom.write_partial(new_image, (<section_name>, ...))
169        ex using explicitly provided map:
170          flashrom.write_partial(new_image, layout_map_all, ('all',))
171    """
172
173    def __init__(self, os_if, keep_temp_files=False, target_is_ec=False):
174        """ constructor of flashrom_util. help(flashrom_util) for more info
175
176        @param os_if: an object providing interface to OS services
177        @param keep_temp_files: if true, preserve temp files after operations
178        @param target_is_ec: if false, target is BIOS/AP
179
180        @type os_if: client.cros.faft.utils.os_interface.OSInterface
181        @type keep_temp_files: bool
182        @type target_is_ec: bool
183        """
184
185        self.os_if = os_if
186        self.keep_temp_files = keep_temp_files
187        self.firmware_layout = {}
188        self._target_command = ''
189        if target_is_ec:
190            self._enable_ec_access()
191        else:
192            self._enable_bios_access()
193
194    def _enable_bios_access(self):
195        if self.os_if.test_mode or self.os_if.target_hosted():
196            self._target_command = '-p host'
197
198    def _enable_ec_access(self):
199        if self.os_if.test_mode or self.os_if.target_hosted():
200            self._target_command = '-p ec'
201
202    def _get_temp_filename(self, prefix):
203        """Returns name of a temporary file in /tmp."""
204        return self.os_if.create_temp_file(prefix)
205
206    def _remove_temp_file(self, filename):
207        """Removes a temp file if self.keep_temp_files is false."""
208        if self.keep_temp_files:
209            return
210        if self.os_if.path_exists(filename):
211            self.os_if.remove_file(filename)
212
213    def _create_layout_file(self, layout_map):
214        """Creates a layout file based on layout_map.
215
216        Returns the file name containing layout information.
217        """
218        layout_text = [
219                '0x%08lX:0x%08lX %s' % (v[0], v[1], k)
220                for k, v in layout_map.items()
221        ]
222        layout_text.sort()  # XXX unstable if range exceeds 2^32
223        tmpfn = self._get_temp_filename('lay_')
224        self.os_if.write_file(tmpfn, '\n'.join(layout_text) + '\n')
225        return tmpfn
226
227    def check_target(self):
228        """Check if flashrom programmer is working, by specifying no commands.
229
230        The command executed is just 'flashrom -p <target>'.
231
232        @return: True if flashrom completed successfully
233        @raise autotest_lib.client.common_lib.error.CmdError: if flashrom failed
234        """
235        cmd = 'flashrom %s' % self._target_command
236        self.os_if.run_shell_command(cmd)
237        return True
238
239    def get_section(self, base_image, section_name):
240        """
241        Retrieves a section of data based on section_name in layout_map.
242        Raises error if unknown section or invalid layout_map.
243        """
244        if section_name not in self.firmware_layout:
245            return ''
246        pos = self.firmware_layout[section_name]
247        if pos[0] >= pos[1] or pos[1] >= len(base_image):
248            raise TestError(
249                    'INTERNAL ERROR: invalid layout map: %s.' % section_name)
250        blob = base_image[pos[0]:pos[1] + 1]
251        # Trim down the main firmware body to its actual size since the
252        # signing utility uses the size of the input file as the size of
253        # the data to sign. Make it the same way as firmware creation.
254        if section_name in ('FVMAIN', 'FVMAINB', 'ECMAINA', 'ECMAINB'):
255            align = 4
256            pad = blob[-1]
257            blob = blob.rstrip(pad)
258            blob = blob + ((align - 1) - (len(blob) - 1) % align) * pad
259        return blob
260
261    def put_section(self, base_image, section_name, data):
262        """
263        Updates a section of data based on section_name in firmware_layout.
264        Raises error if unknown section.
265        Returns the full updated image data.
266        """
267        pos = self.firmware_layout[section_name]
268        if pos[0] >= pos[1] or pos[1] >= len(base_image):
269            raise TestError('INTERNAL ERROR: invalid layout map.')
270        if len(data) != pos[1] - pos[0] + 1:
271            # Pad the main firmware body since we trimed it down before.
272            if (len(data) < pos[1] - pos[0] + 1
273                        and section_name in ('FVMAIN', 'FVMAINB', 'ECMAINA',
274                                             'ECMAINB', 'RW_FWID')):
275                pad = base_image[pos[1]]
276                data = data + pad * (pos[1] - pos[0] + 1 - len(data))
277            else:
278                raise TestError('INTERNAL ERROR: unmatched data size.')
279        return base_image[0:pos[0]] + data + base_image[pos[1] + 1:]
280
281    def get_size(self):
282        """ Gets size of current flash ROM """
283        # TODO(hungte) Newer version of tool (flashrom) may support --get-size
284        # command which is faster in future. Right now we use back-compatible
285        # method: read whole and then get length.
286        image = self.read_whole()
287        return len(image)
288
289    def set_firmware_layout(self, file_name):
290        """get layout read from the BIOS """
291
292        scraper = LayoutScraper(self.os_if)
293        self.firmware_layout = scraper.get_layout(file_name)
294
295    def enable_write_protect(self):
296        """Enable the write protection of the flash chip."""
297
298        # For MTD devices, this will fail: need both --wp-range and --wp-enable.
299        # See: https://crrev.com/c/275381
300
301        cmd = 'flashrom %s --verbose --wp-enable' % self._target_command
302        self.os_if.run_shell_command(cmd, modifies_device=True)
303
304    def disable_write_protect(self):
305        """Disable the write protection of the flash chip."""
306        cmd = 'flashrom %s --verbose --wp-disable' % self._target_command
307        self.os_if.run_shell_command(cmd, modifies_device=True)
308
309    def set_write_protect_region(self, image_file, region, enabled=None):
310        """
311        Set write protection region, using specified image's layout.
312
313        The name should match those seen in `futility dump_fmap <image>`, and
314        is not checked against self.firmware_layout, due to different naming.
315
316        @param image_file: path of the image file to read regions from
317        @param region: Region to set (usually WP_RO)
318        @param enabled: if True, run --wp-enable; if False, run --wp-disable.
319        """
320        cmd = 'flashrom %s --verbose --image %s:%s --wp-region %s' % (
321                self._target_command, region, image_file, region)
322        if enabled is not None:
323            cmd += ' '
324            cmd += '--wp-enable' if enabled else '--wp-disable'
325
326        self.os_if.run_shell_command(cmd, modifies_device=True)
327
328    def set_write_protect_range(self, start, length, enabled=None):
329        """
330        Set write protection range by offset, using current image's layout.
331
332        @param start: offset (bytes) from start of flash to start of range
333        @param length: offset (bytes) from start of range to end of range
334        @param enabled: If True, run --wp-enable; if False, run --wp-disable.
335                        If None (default), don't specify either one.
336        """
337        cmd = 'flashrom %s --verbose --wp-range %s %s' % (
338                self._target_command, start, length)
339        if enabled is not None:
340            cmd += ' '
341            cmd += '--wp-enable' if enabled else '--wp-disable'
342
343        self.os_if.run_shell_command(cmd, modifies_device=True)
344
345    def get_write_protect_status(self):
346        """Get a dict describing the status of the write protection
347
348        @return: {'enabled': True/False, 'start': '0x0', 'length': '0x0', ...}
349        @rtype: dict
350        """
351        # https://crrev.com/8ebbd500b5d8da9f6c1b9b44b645f99352ef62b4/writeprotect.c
352
353        status_pattern = re.compile(
354                r'WP: status: (.*)')
355        enabled_pattern = re.compile(
356                r'WP: write protect is (\w+)\.?')
357        range_pattern = re.compile(
358                r'WP: write protect range: start=(\w+), len=(\w+)')
359        range_err_pattern = re.compile(
360                r'WP: write protect range: (.+)')
361
362        output = self.os_if.run_shell_command_get_output(
363                'flashrom %s --wp-status' % self._target_command)
364
365        wp_status = {}
366        for line in output:
367            if not line.startswith('WP: '):
368                continue
369
370            found_enabled = re.match(enabled_pattern, line)
371            if found_enabled:
372                status_word = found_enabled.group(1)
373                wp_status['enabled'] = (status_word == 'enabled')
374                continue
375
376            found_range = re.match(range_pattern, line)
377            if found_range:
378                (start, length) = found_range.groups()
379                wp_status['start'] = int(start, 16)
380                wp_status['length'] = int(length, 16)
381                continue
382
383            found_range_err = re.match(range_err_pattern, line)
384            if found_range_err:
385                # WP: write protect range: (cannot resolve the range)
386                wp_status['error'] = found_range_err.group(1)
387                continue
388
389            found_status = re.match(status_pattern, line)
390            if found_status:
391                wp_status['status'] = found_status.group(1)
392                continue
393
394        return wp_status
395
396    def dump_flash(self, filename):
397        """Read the flash device's data into a file, but don't parse it."""
398        cmd = 'flashrom %s -r "%s"' % (self._target_command, filename)
399        self.os_if.log('flashrom_util.dump_flash(): %s' % cmd)
400        self.os_if.run_shell_command(cmd)
401
402    def read_whole(self):
403        """
404        Reads whole flash ROM data.
405        Returns the data read from flash ROM, or empty string for other error.
406        """
407        tmpfn = self._get_temp_filename('rd_')
408        cmd = 'flashrom %s -r "%s"' % (self._target_command, tmpfn)
409        self.os_if.log('flashrom_util.read_whole(): %s' % cmd)
410        self.os_if.run_shell_command(cmd)
411        result = self.os_if.read_file(tmpfn)
412        self.set_firmware_layout(tmpfn)
413
414        # clean temporary resources
415        self._remove_temp_file(tmpfn)
416        return result
417
418    def write_partial(self, base_image, write_list, write_layout_map=None):
419        """
420        Writes data in sections of write_list to flash ROM.
421        An exception is raised if write operation fails.
422        """
423
424        if write_layout_map:
425            layout_map = write_layout_map
426        else:
427            layout_map = self.firmware_layout
428
429        tmpfn = self._get_temp_filename('wr_')
430        self.os_if.write_file(tmpfn, base_image)
431        layout_fn = self._create_layout_file(layout_map)
432
433        write_cmd = 'flashrom %s -l "%s" -i %s -w "%s"' % (
434                self._target_command, layout_fn, ' -i '.join(write_list),
435                tmpfn)
436        self.os_if.log('flashrom.write_partial(): %s' % write_cmd)
437        self.os_if.run_shell_command(write_cmd, modifies_device=True)
438
439        # clean temporary resources
440        self._remove_temp_file(tmpfn)
441        self._remove_temp_file(layout_fn)
442
443    def write_whole(self, base_image):
444        """Write the whole base image. """
445        layout_map = {'all': (0, len(base_image) - 1)}
446        self.write_partial(base_image, ('all', ), layout_map)
447
448    def get_write_cmd(self, image=None):
449        """Get the command to write the whole image (no layout handling)
450
451        @param image: the filename (empty to use current handler data)
452        """
453        return 'flashrom %s -w "%s"' % (self._target_command, image)
454