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
5""" This module provides convenience routines to access Flash ROM (EEPROM)
6
7saft_flashrom_util is based on utility 'flashrom'.
8
9Original tool syntax:
10    (read ) flashrom -r <file>
11    (write) flashrom -l <layout_fn> [-i <image_name> ...] -w <file>
12
13The layout_fn is in format of
14    address_begin:address_end image_name
15    which defines a region between (address_begin, address_end) and can
16    be accessed by the name image_name.
17
18Currently the tool supports multiple partial write but not partial read.
19
20In the saft_flashrom_util, we provide read and partial write abilities.
21For more information, see help(saft_flashrom_util.flashrom_util).
22"""
23
24class TestError(Exception):
25    pass
26
27
28class LayoutScraper(object):
29    """Object of this class is used to retrieve layout from a BIOS file."""
30
31    # The default conversion table for mosys.
32    DEFAULT_CHROMEOS_FMAP_CONVERSION = {
33        "Boot Stub": "FV_BSTUB",
34        "GBB Area": "FV_GBB",
35        "Recovery Firmware": "FVDEV",
36        "RO VPD": "RO_VPD",
37        "Firmware A Key": "VBOOTA",
38        "Firmware A Data": "FVMAIN",
39        "Firmware B Key": "VBOOTB",
40        "Firmware B Data": "FVMAINB",
41        "Log Volume": "FV_LOG",
42        # New layout in Chrome OS Main Processor Firmware Specification,
43        # used by all newer (>2011) platforms except Mario.
44        "BOOT_STUB": "FV_BSTUB",
45        "RO_FRID": "RO_FRID",
46        "GBB": "FV_GBB",
47        "RECOVERY": "FVDEV",
48        "VBLOCK_A": "VBOOTA",
49        "VBLOCK_B": "VBOOTB",
50        "FW_MAIN_A": "FVMAIN",
51        "FW_MAIN_B": "FVMAINB",
52        "RW_FWID_A": "RW_FWID_A",
53        "RW_FWID_B": "RW_FWID_B",
54        # Memory Training data cache for recovery boots
55        # Added on Nov 09, 2016
56        "RECOVERY_MRC_CACHE": "RECOVERY_MRC_CACHE",
57        # New sections in Depthcharge.
58        "EC_MAIN_A": "ECMAINA",
59        "EC_MAIN_B": "ECMAINB",
60        # EC firmware layout
61        "EC_RW": "EC_RW",
62        "EC_RW_B": "EC_RW_B",
63        "RW_FWID": "RW_FWID",
64        }
65
66    def __init__(self, os_if):
67        self.image = None
68        self.os_if = os_if
69
70    def _get_text_layout(self, file_name):
71        """Retrieve text layout from a firmware image file.
72
73        This function uses the 'mosys' utility to scan the firmware image and
74        retrieve the section layout information.
75
76        The layout is reported as a set of lines with multiple
77        "<name>"="value" pairs, all this output is passed to the caller.
78        """
79
80        mosys_cmd = 'mosys -k eeprom map %s' % file_name
81        return self.os_if.run_shell_command_get_output(mosys_cmd)
82
83    def _line_to_dictionary(self, line):
84        """Convert a text layout line into a dictionary.
85
86        Get a string consisting of single space separated "<name>"="value>"
87        pairs and convert it into a dictionary where keys are the <name>
88        fields, and values are the corresponding <value> fields.
89
90        Return the dictionary to the caller.
91        """
92
93        rv = {}
94
95        items = line.replace('" ', '"^').split('^')
96        for item in items:
97            pieces = item.split('=')
98            if len(pieces) != 2:
99                continue
100            rv[pieces[0]] = pieces[1].strip('"')
101        return rv
102
103    def check_layout(self, layout, file_size):
104        """Verify the layout to be consistent.
105
106        The layout is consistent if there is no overlapping sections and the
107        section boundaries do not exceed the file size.
108
109        Inputs:
110          layout: a dictionary keyed by a string (the section name) with
111                  values being two integers tuples, the first and the last
112                  bites' offset in the file.
113          file_size: and integer, the size of the file the layout describes
114                     the sections in.
115
116        Raises:
117          TestError in case the layout is not consistent.
118        """
119
120        # Generate a list of section range tuples.
121        ost = sorted([layout[section] for section in layout])
122        base = -1
123        for section_base, section_end in ost:
124            if section_base <= base or section_end + 1 < section_base:
125                # Overlapped section is possible, like the fwid which is
126                # inside the main fw section.
127                self.os_if.log('overlapped section at 0x%x..0x%x' % (
128                        section_base, section_end))
129            base = section_end
130        if base > file_size:
131            raise TestError('Section end 0x%x exceeds file size %x' % (
132                    base, file_size))
133
134    def get_layout(self, file_name):
135        """Generate layout for a firmware file.
136
137        First retrieve the text layout as reported by 'mosys' and then convert
138        it into a dictionary, replacing section names reported by mosys into
139        matching names from DEFAULT_CHROMEOS_FMAP_CONVERSION dictionary above,
140        using the names as keys in the layout dictionary. The elements of the
141        layout dictionary are the offsets of the first ans last bytes of the
142        section in the firmware file.
143
144        Then verify the generated layout's consistency and return it to the
145        caller.
146        """
147
148        layout_data = {} # keyed by the section name, elements - tuples of
149                         # (<section start addr>, <section end addr>)
150
151        for line in self._get_text_layout(file_name):
152            d = self._line_to_dictionary(line)
153            try:
154                name = self.DEFAULT_CHROMEOS_FMAP_CONVERSION[d['area_name']]
155            except KeyError:
156                continue  # This line does not contain an area of interest.
157
158            if name in layout_data:
159                raise TestError('%s duplicated in the layout' % name)
160
161            offset = int(d['area_offset'], 0)
162            size = int(d['area_size'], 0)
163            layout_data[name] = (offset, offset + size - 1)
164
165        self.check_layout(layout_data, self.os_if.get_file_size(file_name))
166        return layout_data
167
168# flashrom utility wrapper
169class flashrom_util(object):
170    """ a wrapper for "flashrom" utility.
171
172    You can read, write, or query flash ROM size with this utility.
173    Although you can do "partial-write", the tools always takes a
174    full ROM image as input parameter.
175
176    NOTE before accessing flash ROM, you may need to first "select"
177    your target - usually BIOS or EC. That part is not handled by
178    this utility. Please find other external script to do it.
179
180    To perform a read, you need to:
181     1. Prepare a flashrom_util object
182        ex: flashrom = flashrom_util.flashrom_util()
183     2. Perform read operation
184        ex: image = flashrom.read_whole()
185
186        When the contents of the flashrom is read off the target, it's map
187        gets created automatically (read from the flashrom image using
188        'mosys'). If the user wants this object to operate on some other file,
189        he could either have the map for the file created explicitly by
190        invoking flashrom.set_firmware_layout(filename), or supply his own map
191        (which is a dictionary where keys are section names, and values are
192        tuples of integers, base address of the section and the last address
193        of the section).
194
195    By default this object operates on the map retrieved from the image and
196    stored locally, this map can be overwritten by an explicitly passed user
197    map.
198
199   To perform a (partial) write:
200
201     1. Prepare a buffer storing an image to be written into the flashrom.
202     2. Have the map generated automatically or prepare your own, for instance:
203        ex: layout_map_all = { 'all': (0, rom_size - 1) }
204        ex: layout_map = { 'ro': (0, 0xFFF), 'rw': (0x1000, rom_size-1) }
205     4. Perform write operation
206
207        ex using default map:
208          flashrom.write_partial(new_image, (<section_name>, ...))
209        ex using explicitly provided map:
210          flashrom.write_partial(new_image, layout_map_all, ('all',))
211
212    Attributes:
213        keep_temp_files: boolean flag to control cleaning of temporary files
214    """
215
216    def __init__(self, os_if, keep_temp_files=False,
217                 target_is_ec=False):
218        """ constructor of flashrom_util. help(flashrom_util) for more info """
219        self.os_if = os_if
220        self.keep_temp_files = keep_temp_files
221        self.firmware_layout = {}
222        self._target_command = ''
223        if target_is_ec:
224            self._enable_ec_access()
225        else:
226            self._enable_bios_access()
227
228    def _enable_bios_access(self):
229        if not self.os_if.target_hosted():
230            return
231        self._target_command = '-p host'
232
233    def _enable_ec_access(self):
234        if not self.os_if.target_hosted():
235            return
236        self._target_command = '-p ec'
237
238    def _get_temp_filename(self, prefix):
239        """Returns name of a temporary file in /tmp."""
240        return self.os_if.create_temp_file(prefix)
241
242    def _remove_temp_file(self, filename):
243        """Removes a temp file if self.keep_temp_files is false."""
244        if self.keep_temp_files:
245            return
246        if self.os_if.path_exists(filename):
247            self.os_if.remove_file(filename)
248
249    def _create_layout_file(self, layout_map):
250        """Creates a layout file based on layout_map.
251
252        Returns the file name containing layout information.
253        """
254        layout_text = ['0x%08lX:0x%08lX %s' % (v[0], v[1], k)
255            for k, v in layout_map.items()]
256        layout_text.sort()  # XXX unstable if range exceeds 2^32
257        tmpfn = self._get_temp_filename('lay_')
258        self.os_if.write_file(tmpfn, '\n'.join(layout_text) + '\n')
259        return tmpfn
260
261    def get_section(self, base_image, section_name):
262        """
263        Retrieves a section of data based on section_name in layout_map.
264        Raises error if unknown section or invalid layout_map.
265        """
266        if section_name not in self.firmware_layout:
267            return ''
268        pos = self.firmware_layout[section_name]
269        if pos[0] >= pos[1] or pos[1] >= len(base_image):
270            raise TestError('INTERNAL ERROR: invalid layout map: %s.' %
271                            section_name)
272        blob = base_image[pos[0] : pos[1] + 1]
273        # Trim down the main firmware body to its actual size since the
274        # signing utility uses the size of the input file as the size of
275        # the data to sign. Make it the same way as firmware creation.
276        if section_name in ('FVMAIN', 'FVMAINB', 'ECMAINA', 'ECMAINB'):
277            align = 4
278            pad = blob[-1]
279            blob = blob.rstrip(pad)
280            blob = blob + ((align - 1) - (len(blob) - 1) % align) * pad
281        return blob
282
283    def put_section(self, base_image, section_name, data):
284        """
285        Updates a section of data based on section_name in firmware_layout.
286        Raises error if unknown section.
287        Returns the full updated image data.
288        """
289        pos = self.firmware_layout[section_name]
290        if pos[0] >= pos[1] or pos[1] >= len(base_image):
291            raise TestError('INTERNAL ERROR: invalid layout map.')
292        if len(data) != pos[1] - pos[0] + 1:
293            # Pad the main firmware body since we trimed it down before.
294            if (len(data) < pos[1] - pos[0] + 1 and section_name in
295                    ('FVMAIN', 'FVMAINB', 'ECMAINA', 'ECMAINB',
296                     'RW_FWID')):
297                pad = base_image[pos[1]]
298                data = data + pad * (pos[1] - pos[0] + 1 - len(data))
299            else:
300                raise TestError('INTERNAL ERROR: unmatched data size.')
301        return base_image[0 : pos[0]] + data + base_image[pos[1] + 1 :]
302
303    def get_size(self):
304        """ Gets size of current flash ROM """
305        # TODO(hungte) Newer version of tool (flashrom) may support --get-size
306        # command which is faster in future. Right now we use back-compatible
307        # method: read whole and then get length.
308        image = self.read_whole()
309        return len(image)
310
311    def set_firmware_layout(self, file_name):
312        """get layout read from the BIOS """
313
314        scraper = LayoutScraper(self.os_if)
315        self.firmware_layout = scraper.get_layout(file_name)
316
317    def enable_write_protect(self):
318        """Enable the write pretection of the flash chip."""
319        cmd = 'flashrom %s --wp-enable' % self._target_command
320        self.os_if.run_shell_command(cmd)
321
322    def disable_write_protect(self):
323        """Disable the write pretection of the flash chip."""
324        cmd = 'flashrom %s --wp-disable' % self._target_command
325        self.os_if.run_shell_command(cmd)
326
327    def read_whole(self):
328        """
329        Reads whole flash ROM data.
330        Returns the data read from flash ROM, or empty string for other error.
331        """
332        tmpfn = self._get_temp_filename('rd_')
333        cmd = 'flashrom %s -r "%s"' % (self._target_command, tmpfn)
334        self.os_if.log('flashrom_util.read_whole(): %s' % cmd)
335        self.os_if.run_shell_command(cmd)
336        result = self.os_if.read_file(tmpfn)
337        self.set_firmware_layout(tmpfn)
338
339        # clean temporary resources
340        self._remove_temp_file(tmpfn)
341        return result
342
343    def write_partial(self, base_image, write_list, write_layout_map=None):
344        """
345        Writes data in sections of write_list to flash ROM.
346        An exception is raised if write operation fails.
347        """
348
349        if write_layout_map:
350            layout_map = write_layout_map
351        else:
352            layout_map = self.firmware_layout
353
354        tmpfn = self._get_temp_filename('wr_')
355        self.os_if.write_file(tmpfn, base_image)
356        layout_fn = self._create_layout_file(layout_map)
357
358        cmd = 'flashrom %s -l "%s" -i %s -w "%s"' % (
359                self._target_command, layout_fn, ' -i '.join(write_list), tmpfn)
360        self.os_if.log('flashrom.write_partial(): %s' % cmd)
361        self.os_if.run_shell_command(cmd)
362
363        # flashrom write will reboot the ec after corruption
364        # For Android, need to make sure ec is back online
365        # before continuing, or adb command will cause test failure
366        if self.os_if.is_android:
367            self.os_if.wait_for_device(60)
368
369        # clean temporary resources
370        self._remove_temp_file(tmpfn)
371        self._remove_temp_file(layout_fn)
372
373    def write_whole(self, base_image):
374        """Write the whole base image. """
375        layout_map = { 'all': (0, len(base_image) - 1) }
376        self.write_partial(base_image, ('all',), layout_map)
377