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