flash_algo.py

00001 #!/usr/bin/env python
00002 """
00003  mbed
00004  Copyright (c) 2017-2017 ARM Limited
00005 
00006  Licensed under the Apache License, Version 2.0 (the "License");
00007  you may not use this file except in compliance with the License.
00008  You may obtain a copy of the License at
00009 
00010      http://www.apache.org/licenses/LICENSE-2.0
00011 
00012  Unless required by applicable law or agreed to in writing, software
00013  distributed under the License is distributed on an "AS IS" BASIS,
00014  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
00015  See the License for the specific language governing permissions and
00016  limitations under the License.
00017 """
00018 
00019 from __future__ import print_function
00020 import os
00021 import struct
00022 import binascii
00023 import argparse
00024 import logging
00025 import StringIO
00026 import jinja2
00027 from collections import namedtuple
00028 from itertools import count
00029 
00030 from elftools.common.py3compat import bytes2str
00031 from elftools.elf.elffile import ELFFile
00032 from elftools.elf.sections import SymbolTableSection
00033 
00034 logger = logging.getLogger(__name__)
00035 logger.addHandler(logging.NullHandler())
00036 
00037 
00038 def main():
00039     parser = argparse.ArgumentParser(description="Algo Extracter")
00040     parser.add_argument("input", help="File to extract flash algo from")
00041     parser.add_argument("template", default="py_blob.tmpl",
00042                         help="Template to use")
00043     parser.add_argument("output", help="Output file")
00044     args = parser.parse_args()
00045 
00046     with open(args.input, "rb") as file_handle:
00047         data = file_handle.read()
00048     algo = PackFlashAlgo(data)
00049     algo.process_template(args.template, args.output)
00050 
00051 
00052 class PackFlashAlgo (object):
00053     """
00054     Class to wrap a flash algo
00055 
00056     This class is intended to provide easy access to the information
00057     provided by a flash algorithm, such as symbols and the flash
00058     algorithm itself.
00059     """
00060 
00061     REQUIRED_SYMBOLS = set([
00062         "Init",
00063         "UnInit",
00064         "EraseSector",
00065         "ProgramPage",
00066     ])
00067 
00068     EXTRA_SYMBOLS = set([
00069         "BlankCheck",
00070         "EraseChip",
00071         "Verify",
00072     ])
00073 
00074     def __init__ (self, data):
00075         """Construct a PackFlashAlgorithm from an ElfFileSimple"""
00076         self.elf  = ElfFileSimple(data)
00077         self.flash_info  = PackFlashInfo(self.elf )
00078 
00079         self.flash_start  = self.flash_info .start
00080         self.flash_size  = self.flash_info .size
00081         self.page_size  = self.flash_info .page_size
00082         self.sector_sizes  = self.flash_info .sector_info_list
00083 
00084         symbols = {}
00085         symbols.update(_extract_symbols(self.elf , self.REQUIRED_SYMBOLS ))
00086         symbols.update(_extract_symbols(self.elf , self.EXTRA_SYMBOLS ,
00087                                         default=0xFFFFFFFF))
00088         self.symbols  = symbols
00089 
00090         sections_to_find = (
00091             ("PrgCode", "SHT_PROGBITS"),
00092             ("PrgData", "SHT_PROGBITS"),
00093             ("PrgData", "SHT_NOBITS"),
00094         )
00095 
00096         ro_rw_zi = _find_sections(self.elf , sections_to_find)
00097         ro_rw_zi = _algo_fill_zi_if_missing(ro_rw_zi)
00098         error_msg = _algo_check_for_section_problems(ro_rw_zi)
00099         if error_msg is not None:
00100             raise Exception(error_msg)
00101 
00102         sect_ro, sect_rw, sect_zi = ro_rw_zi
00103         self.ro_start  = sect_ro["sh_addr"]
00104         self.ro_size  = sect_ro["sh_size"]
00105         self.rw_start  = sect_rw["sh_addr"]
00106         self.rw_size  = sect_rw["sh_size"]
00107         self.zi_start  = sect_zi["sh_addr"]
00108         self.zi_size  = sect_zi["sh_size"]
00109 
00110         self.algo_data  = _create_algo_bin(ro_rw_zi)
00111 
00112     def format_algo_data (self, spaces, group_size, fmt):
00113         """"
00114         Return a string representing algo_data suitable for use in a template
00115 
00116         The string is intended for use in a template.
00117 
00118         :param spaces: The number of leading spaces for each line
00119         :param group_size: number of elements per line (element type
00120             depends of format)
00121         :param fmt: - format to create - can be either "hex" or "c"
00122         """
00123         padding = " " * spaces
00124         if fmt == "hex":
00125             blob = binascii.b2a_hex(self.algo_data )
00126             line_list = []
00127             for i in xrange(0, len(blob), group_size):
00128                 line_list.append('"' + blob[i:i + group_size] + '"')
00129             return ("\n" + padding).join(line_list)
00130         elif fmt == "c":
00131             blob = self.algo_data [:]
00132             pad_size = 0 if len(blob) % 4 == 0 else 4 - len(blob) % 4
00133             blob = blob + "\x00" * pad_size
00134             integer_list = struct.unpack("<" + "L" * (len(blob) / 4), blob)
00135             line_list = []
00136             for pos in range(0, len(integer_list), group_size):
00137                 group = ["0x%08x" % value for value in
00138                          integer_list[pos:pos + group_size]]
00139                 line_list.append(", ".join(group))
00140             return (",\n" + padding).join(line_list)
00141         else:
00142             raise Exception("Unsupported format %s" % fmt)
00143 
00144     def process_template (self, template_path, output_path, data_dict=None):
00145         """
00146         Generate output from the supplied template
00147 
00148         All the public methods and fields of this class can be accessed from
00149         the template via "algo".
00150 
00151         :param template_path: Relative or absolute file path to the template
00152         :param output_path: Relative or absolute file path to create
00153         :param data_dict: Additional data to use when generating
00154         """
00155         if data_dict is None:
00156             data_dict = {}
00157         else:
00158             assert isinstance(data_dict, dict)
00159             data_dict = dict(data_dict)
00160         assert "algo" not in data_dict, "algo already set by user data"
00161         data_dict["algo"] = self
00162 
00163         with open(template_path) as file_handle:
00164             template_text = file_handle.read()
00165 
00166         template = jinja2.Template(template_text)
00167         target_text = template.render(data_dict)
00168 
00169         with open(output_path, "wb") as file_handle:
00170             file_handle.write(target_text)
00171 
00172 
00173 def _extract_symbols(simple_elf, symbols, default=None):
00174     """Fill 'symbols' field with required flash algo symbols"""
00175     to_ret = {}
00176     for symbol in symbols:
00177         if symbol not in simple_elf.symbols:
00178             if default is not None:
00179                 to_ret[symbol] = default
00180                 continue
00181             raise Exception("Missing symbol %s" % symbol)
00182         to_ret[symbol] = simple_elf.symbols[symbol].value
00183     return to_ret
00184 
00185 
00186 def _find_sections(elf, name_type_pairs):
00187     """Return a list of sections the same length and order of the input list"""
00188     sections = [None] * len(name_type_pairs)
00189     for section in elf.iter_sections():
00190         section_name = bytes2str(section.name)
00191         section_type = section["sh_type"]
00192         for i, name_and_type in enumerate(name_type_pairs):
00193             if name_and_type != (section_name, section_type):
00194                 continue
00195             if sections[i] is not None:
00196                 raise Exception("Elf contains duplicate section %s attr %s" %
00197                                 (section_name, section_type))
00198             sections[i] = section
00199     return sections
00200 
00201 
00202 def _algo_fill_zi_if_missing(ro_rw_zi):
00203     """Create an empty zi section if it is missing"""
00204     s_ro, s_rw, s_zi = ro_rw_zi
00205     if s_rw is None:
00206         return ro_rw_zi
00207     if s_zi is not None:
00208         return ro_rw_zi
00209     s_zi = {
00210         "sh_addr": s_rw["sh_addr"] + s_rw["sh_size"],
00211         "sh_size": 0
00212     }
00213     return s_ro, s_rw, s_zi
00214 
00215 
00216 def _algo_check_for_section_problems(ro_rw_zi):
00217     """Return a string describing any errors with the layout or None if good"""
00218     s_ro, s_rw, s_zi = ro_rw_zi
00219     if s_ro is None:
00220         return "RO section is missing"
00221     if s_rw is None:
00222         return "RW section is missing"
00223     if s_zi is None:
00224         return "ZI section is missing"
00225     if s_ro["sh_addr"] != 0:
00226         return "RO section does not start at address 0"
00227     if s_ro["sh_addr"] + s_ro["sh_size"] != s_rw["sh_addr"]:
00228         return "RW section does not follow RO section"
00229     if s_rw["sh_addr"] + s_rw["sh_size"] != s_zi["sh_addr"]:
00230         return "ZI section does not follow RW section"
00231     return None
00232 
00233 
00234 def _create_algo_bin(ro_rw_zi):
00235     """Create a binary blob of the flash algo which can execute from ram"""
00236     sect_ro, sect_rw, sect_zi = ro_rw_zi
00237     algo_size = sect_ro["sh_size"] + sect_rw["sh_size"] + sect_zi["sh_size"]
00238     algo_data = bytearray(algo_size)
00239     for section in (sect_ro, sect_rw):
00240         start = section["sh_addr"]
00241         size = section["sh_size"]
00242         data = section.data()
00243         assert len(data) == size
00244         algo_data[start:start + size] = data
00245     return algo_data
00246 
00247 
00248 class PackFlashInfo (object):
00249     """Wrapper class for the non-executable information in an FLM file"""
00250 
00251     FLASH_DEVICE_STRUCT = "<H128sHLLLLBxxxLL"
00252     FLASH_SECTORS_STRUCT = "<LL"
00253     FLASH_SECTORS_STRUCT_SIZE = struct.calcsize(FLASH_SECTORS_STRUCT)
00254     SECTOR_END = 0xFFFFFFFF
00255 
00256     def __init__(self, elf_simple):
00257         dev_info = elf_simple.symbols["FlashDevice"]
00258         info_start = dev_info.value
00259         info_size = struct.calcsize(self.FLASH_DEVICE_STRUCT )
00260         data = elf_simple.read(info_start, info_size)
00261         values = struct.unpack(self.FLASH_DEVICE_STRUCT , data)
00262 
00263         self.version  = values[0]
00264         self.name  = values[1].strip("\x00")
00265         self.type  = values[2]
00266         self.start  = values[3]
00267         self.size  = values[4]
00268         self.page_size  = values[5]
00269         self.value_empty  = values[7]
00270         self.prog_timeout_ms  = values[8]
00271         self.erase_timeout_ms  = values[9]
00272 
00273         sector_gen = self._sector_and_sz_itr (elf_simple,
00274                                              info_start + info_size)
00275         self.sector_info_list  = list(sector_gen)
00276 
00277     def __str__(self):
00278         desc = ""
00279         desc += "Flash Device:" + os.linesep
00280         desc += "  name=%s" % self.name  + os.linesep
00281         desc += "  version=0x%x" % self.version  + os.linesep
00282         desc += "  type=%i" % self.type  + os.linesep
00283         desc += "  start=0x%x" % self.start  + os.linesep
00284         desc += "  size=0x%x" % self.size  + os.linesep
00285         desc += "  page_size=0x%x" % self.page_size  + os.linesep
00286         desc += "  value_empty=0x%x" % self.value_empty  + os.linesep
00287         desc += "  prog_timeout_ms=%i" % self.prog_timeout_ms  + os.linesep
00288         desc += "  erase_timeout_ms=%i" % self.erase_timeout_ms  + os.linesep
00289         desc += "  sectors:" + os.linesep
00290         for sector_start, sector_size in self.sector_info_list :
00291             desc += ("    start=0x%x, size=0x%x" %
00292                      (sector_start, sector_size) + os.linesep)
00293         return desc
00294 
00295     def _sector_and_sz_itr(self, elf_simple, data_start):
00296         """Iterator which returns starting address and sector size"""
00297         for entry_start in count(data_start, self.FLASH_SECTORS_STRUCT_SIZE ):
00298             data = elf_simple.read(entry_start, self.FLASH_SECTORS_STRUCT_SIZE )
00299             size, start = struct.unpack(self.FLASH_SECTORS_STRUCT , data)
00300             start_and_size = start, size
00301             if start_and_size == (self.SECTOR_END , self.SECTOR_END ):
00302                 return
00303             yield start_and_size
00304 
00305 
00306 SymbolSimple = namedtuple("SymbolSimple", "name, value, size")
00307 
00308 
00309 class ElfFileSimple (ELFFile):
00310     """Wrapper for elf object which allows easy access to symbols and rom"""
00311 
00312     def __init__ (self, data):
00313         """Construct a ElfFileSimple from bytes or a bytearray"""
00314         super(ElfFileSimple, self).__init__(StringIO.StringIO(data))
00315         self.symbols  = self._read_symbol_table ()
00316 
00317     def _read_symbol_table(self):
00318         """Read the symbol table into the field "symbols" for easy use"""
00319         section = self.get_section_by_name(b".symtab")
00320         if not section:
00321             raise Exception("Missing symbol table")
00322 
00323         if not isinstance(section, SymbolTableSection):
00324             raise Exception("Invalid symbol table section")
00325 
00326         symbols = {}
00327         for symbol in section.iter_symbols():
00328             name_str = bytes2str(symbol.name)
00329             if name_str in symbols:
00330                 logging.debug("Duplicate symbol %s", name_str)
00331             symbols[name_str] = SymbolSimple(name_str, symbol["st_value"],
00332                                              symbol["st_size"])
00333         return symbols
00334 
00335     def read (self, addr, size):
00336         """Read program data from the elf file
00337 
00338         :param addr: physical address (load address) to read from
00339         :param size: number of bytes to read
00340         :return: Requested data or None if address is unmapped
00341         """
00342         for segment in self.iter_segments():
00343             seg_addr = segment["p_paddr"]
00344             seg_size = min(segment["p_memsz"], segment["p_filesz"])
00345             if addr >= seg_addr + seg_size:
00346                 continue
00347             if addr + size <= seg_addr:
00348                 continue
00349             # There is at least some overlap
00350 
00351             if addr >= seg_addr and addr + size <= seg_addr + seg_size:
00352                 # Region is fully contained
00353                 data = segment.data()
00354                 start = addr - seg_addr
00355                 return data[start:start + size]
00356 
00357 
00358 if __name__ == '__main__':
00359     main()