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