Make SAFT use layout retrieved from the BIOS image.

saft/flashrom_handler.py

 #!/usr/bin/python
 # Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 '''A module to support automated testing of ChromeOS formware.
 
 Utilizes services provided by <board arch>/flashrom_util.py read/write the
 flashrom chip and to parse the flash rom image.
 
@@ skipping 21 matching lines @@
     '''An object to provide logical services for automated flashrom testing.'''
 
     DELTA = 1  # value to add to a byte to corrupt a section contents
 
     # File in the state directory to store public root key.
     PUB_KEY_FILE_NAME = 'root.pubkey'
 
     def __init__(self):
     # make sure it does not accidentally overwrite the image.
         self.fum = None
-        self.bios_layout = None
         self.chros_if = None
         self.image = ''
         self.pub_key_file = ''
         self.fv_sections = {
             'a': FvImage('VBOOTA', 'FVMAIN'),
             'b': FvImage('VBOOTB', 'FVMAINB'),
             }
 
     def init(self, flashrom_util_module, chros_if, pub_key_file=None):
         '''Flashrom handler initializer.
@@ skipping 17 matching lines @@
                        flashrom is read and its contents are saved into a
                        temporary file which is used instead.
 
         The input file is parsed and the sections of importance (as defined in
         self.fv_sections) are saved in separate files in the state directory
         as defined in the chros_if object.
         '''
 
         if image_file:
             self.image = open(image_file, 'rb').read()
+            self.fum.set_bios_layout(image_file)
         else:
             self.image = self.fum.read_whole()
-        self.bios_layout = self.fum.detect_chromeos_layout('bios',
-                len(self.image))
-        self.whole_flash_layout = self.fum.detect_layout('all',
-                len(self.image))
 
         for section in self.fv_sections.itervalues():
             for subsection_name in section.names():
                 f = open(self.chros_if.state_dir_file(subsection_name), 'wb')
-                f.write(self.fum.get_section(self.image,
-                                             self.bios_layout, subsection_name))
+                f.write(self.fum.get_section(self.image, subsection_name))
                 f.close()
 
         if not self.pub_key_file:
             self._retrieve_pub_key()
 
     def _retrieve_pub_key(self):
         '''Retrieve root public key from the firmware GBB section.'''
 
         gbb_header_format = '<4s20s2I'
         pubk_header_format = '<2Q'
 
-        gbb_section = self.fum.get_section(
-            self.image, self.bios_layout, 'FV_GBB')
+        gbb_section = self.fum.get_section(self.image, 'FV_GBB')
 
         # do some sanity checks
         try:
             sig, _, rootk_offs, rootk_size = struct.unpack_from(
                 gbb_header_format, gbb_section)
         except struct.error, e:
             raise FlashromHandlerError(e)
 
         if sig != '$GBB' or (rootk_offs + rootk_size) > len(gbb_section):
             raise FlashromHandlerError('Bad gbb header')
@@ skipping 49 matching lines @@
 
         if not self.image:
             raise FlashromHandlerError(
                 'Attempt at using an uninitialized object')
         if section not in self.fv_sections:
             raise FlashromHandlerError('Unknown FW section %s'
                                        % section)
 
         # Get the appropriate section of the image.
         subsection_name = self.fv_sections[section].body_name
-        body = self.fum.get_section(self.image, self.bios_layout,
-                                    subsection_name)
+        body = self.fum.get_section(self.image, subsection_name)
 
         # Modify the byte in it within 5% of the section body.
         index = len(body) / 20
         body_list = list(body)
         body_list[index] = '%c' % ((ord(body[index]) + delta) % 0x100)
-        self.image = self.fum.put_section(self.image, self.bios_layout,
+        self.image = self.fum.put_section(self.image,
                                           subsection_name, ''.join(body_list))
         return subsection_name
 
     def corrupt_section(self, section):
         '''Corrupt a section of the image'''
 
         return self._modify_section(section, self.DELTA)
 
     def restore_section(self, section):
         '''Restore a previously corrupted section of the image.'''
 
         return self._modify_section(section, -self.DELTA)
 
     def corrupt_firmware(self, section):
         '''Corrupt a section in the FLASHROM!!!'''
 
         subsection_name = self.corrupt_section(section)
-        self.fum.write_partial(self.image, self.bios_layout,
-                               (subsection_name, ))
+        self.fum.write_partial(self.image, (subsection_name, ))
 
     def restore_firmware(self, section):
         '''Restore the previously corrupted section in the FLASHROM!!!'''
 
         subsection_name = self.restore_section(section)
-        self.fum.write_partial(self.image, self.bios_layout,
-                               (subsection_name, ))
+        self.fum.write_partial(self.image, (subsection_name, ))
 
     def write_whole(self):
         '''Write the whole image into the flashrom.'''
 
         if not self.image:
             raise FlashromHandlerError(
                 'Attempt at using an uninitialized object')
-        self.fum.write_partial(self.image, self.whole_flash_layout, ('all', ))
+        self.fum.write_whole(self.image)
 
     def dump_whole(self, filename):
         '''Write the whole image into a file.'''
 
         if not self.image:
             raise FlashromHandlerError(
                 'Attempt at using an uninitialized object')
         open(filename, 'w').write(self.image)