class RegionMap(object): """ Mostly used in SimAbstractMemory, RegionMap stores a series of mappings between concrete memory address ranges and memory regions, like stack frames and heap regions. """ def __init__(self, is_stack): """ Constructor :param is_stack: Whether this is a region map for stack frames or not. Different strategies apply for stack regions. """ self.is_stack = is_stack # An AVLTree, which maps stack addresses to region IDs self._address_to_region_id = AVLTree() # A dict, which maps region IDs to memory address ranges self._region_id_to_address = {} # # Properties # def __repr__(self): return "RegionMap<%s>" % ("S" if self.is_stack else "H") @property def is_empty(self): return len(self._address_to_region_id) == 0 @property def stack_base(self): if not self.is_stack: raise SimRegionMapError('Calling "stack_base" on a non-stack region map.') return self._address_to_region_id.max_key() @property def region_ids(self): return self._region_id_to_address.keys() # # Public methods # def copy(self): r = RegionMap(is_stack=self.is_stack) # A shallow copy should be enough, since we never modify any RegionDescriptor object in-place if len(self._address_to_region_id) > 0: # TODO: There is a bug in bintrees 2.0.2 that prevents us from copying a non-empty AVLTree object # TODO: Consider submit a pull request r._address_to_region_id = self._address_to_region_id.copy() r._region_id_to_address = self._region_id_to_address.copy() return r def map(self, absolute_address, region_id, related_function_address=None): """ Add a mapping between an absolute address and a region ID. If this is a stack region map, all stack regions beyond (lower than) this newly added regions will be discarded. :param absolute_address: An absolute memory address. :param region_id: ID of the memory region. :param related_function_address: A related function address, mostly used for stack regions. """ if self.is_stack: # Sanity check if not region_id.startswith("stack_"): raise SimRegionMapError('Received a non-stack memory ID "%d" in a stack region map' % region_id) # Remove all stack regions that are lower than the one to add while True: try: addr = self._address_to_region_id.floor_key(absolute_address) descriptor = self._address_to_region_id[addr] # Remove this mapping del self._address_to_region_id[addr] # Remove this region ID from the other mapping del self._region_id_to_address[descriptor.region_id] except KeyError: break else: if absolute_address in self._address_to_region_id: descriptor = self._address_to_region_id[absolute_address] # Remove this mapping del self._address_to_region_id[absolute_address] del self._region_id_to_address[descriptor.region_id] # Add this new region mapping desc = RegionDescriptor(region_id, absolute_address, related_function_address=related_function_address) self._address_to_region_id[absolute_address] = desc self._region_id_to_address[region_id] = desc def unmap_by_address(self, absolute_address): """ Removes a mapping based on its absolute address. :param absolute_address: An absolute address """ desc = self._address_to_region_id[absolute_address] del self._address_to_region_id[absolute_address] del self._region_id_to_address[desc.region_id] def absolutize(self, region_id, relative_address): """ Convert a relative address in some memory region to an absolute address. :param region_id: The memory region ID :param relative_address: The relative memory offset in that memory region :return: An absolute address if converted, or an exception is raised when region id does not exist. """ if region_id == "global": # The global region always bases 0 return relative_address if region_id not in self._region_id_to_address: raise SimRegionMapError('Non-existent region ID "%s"' % region_id) base_address = self._region_id_to_address[region_id].base_address return base_address + relative_address def relativize(self, absolute_address, target_region_id=None): """ Convert an absolute address to the memory offset in a memory region. Note that if an address belongs to heap region is passed in to a stack region map, it will be converted to an offset included in the closest stack frame, and vice versa for passing a stack address to a heap region. Therefore you should only pass in address that belongs to the same category (stack or non-stack) of this region map. :param absolute_address: An absolute memory address :return: A tuple of the closest region ID, the relative offset, and the related function address. """ if target_region_id is None: if self.is_stack: # Get the base address of the stack frame it belongs to base_address = self._address_to_region_id.ceiling_key(absolute_address) else: try: base_address = self._address_to_region_id.floor_key(absolute_address) except KeyError: # Not found. It belongs to the global region then. return "global", absolute_address, None descriptor = self._address_to_region_id[base_address] else: if target_region_id == "global": # Just return the absolute address return "global", absolute_address, None if target_region_id not in self._region_id_to_address: raise SimRegionMapError('Trying to relativize to a non-existent region "%s"' % target_region_id) descriptor = self._region_id_to_address[target_region_id] base_address = descriptor.base_address return descriptor.region_id, absolute_address - base_address, descriptor.related_function_address
class RegionMap(object): """ Mostly used in SimAbstractMemory, RegionMap stores a series of mappings between concrete memory address ranges and memory regions, like stack frames and heap regions. """ def __init__(self, is_stack): """ Constructor :param is_stack: Whether this is a region map for stack frames or not. Different strategies apply for stack regions. """ self.is_stack = is_stack # An AVLTree, which maps stack addresses to region IDs self._address_to_region_id = AVLTree() # A dict, which maps region IDs to memory address ranges self._region_id_to_address = { } # # Properties # def __repr__(self): return "RegionMap<%s>" % ( "S" if self.is_stack else "H" ) @property def is_empty(self): return len(self._address_to_region_id) == 0 @property def stack_base(self): if not self.is_stack: raise SimRegionMapError('Calling "stack_base" on a non-stack region map.') return self._address_to_region_id.max_key() @property def region_ids(self): return self._region_id_to_address.keys() # # Public methods # def copy(self): r = RegionMap(is_stack=self.is_stack) # A shallow copy should be enough, since we never modify any RegionDescriptor object in-place if len(self._address_to_region_id) > 0: # TODO: There is a bug in bintrees 2.0.2 that prevents us from copying a non-empty AVLTree object # TODO: Consider submit a pull request r._address_to_region_id = self._address_to_region_id.copy() r._region_id_to_address = self._region_id_to_address.copy() return r def map(self, absolute_address, region_id, related_function_address=None): """ Add a mapping between an absolute address and a region ID. If this is a stack region map, all stack regions beyond (lower than) this newly added regions will be discarded. :param absolute_address: An absolute memory address. :param region_id: ID of the memory region. :param related_function_address: A related function address, mostly used for stack regions. """ if self.is_stack: # Sanity check if not region_id.startswith('stack_'): raise SimRegionMapError('Received a non-stack memory ID "%d" in a stack region map' % region_id) # Remove all stack regions that are lower than the one to add while True: try: addr = self._address_to_region_id.floor_key(absolute_address) descriptor = self._address_to_region_id[addr] # Remove this mapping del self._address_to_region_id[addr] # Remove this region ID from the other mapping del self._region_id_to_address[descriptor.region_id] except KeyError: break else: if absolute_address in self._address_to_region_id: descriptor = self._address_to_region_id[absolute_address] # Remove this mapping del self._address_to_region_id[absolute_address] del self._region_id_to_address[descriptor.region_id] # Add this new region mapping desc = RegionDescriptor( region_id, absolute_address, related_function_address=related_function_address ) self._address_to_region_id[absolute_address] = desc self._region_id_to_address[region_id] = desc def unmap_by_address(self, absolute_address): """ Removes a mapping based on its absolute address. :param absolute_address: An absolute address """ desc = self._address_to_region_id[absolute_address] del self._address_to_region_id[absolute_address] del self._region_id_to_address[desc.region_id] def absolutize(self, region_id, relative_address): """ Convert a relative address in some memory region to an absolute address. :param region_id: The memory region ID :param relative_address: The relative memory offset in that memory region :return: An absolute address if converted, or an exception is raised when region id does not exist. """ if region_id == 'global': # The global region always bases 0 return relative_address if region_id not in self._region_id_to_address: raise SimRegionMapError('Non-existent region ID "%s"' % region_id) base_address = self._region_id_to_address[region_id].base_address return base_address + relative_address def relativize(self, absolute_address, target_region_id=None): """ Convert an absolute address to the memory offset in a memory region. Note that if an address belongs to heap region is passed in to a stack region map, it will be converted to an offset included in the closest stack frame, and vice versa for passing a stack address to a heap region. Therefore you should only pass in address that belongs to the same category (stack or non-stack) of this region map. :param absolute_address: An absolute memory address :return: A tuple of the closest region ID, the relative offset, and the related function address. """ if target_region_id is None: if self.is_stack: # Get the base address of the stack frame it belongs to base_address = self._address_to_region_id.ceiling_key(absolute_address) else: try: base_address = self._address_to_region_id.floor_key(absolute_address) except KeyError: # Not found. It belongs to the global region then. return 'global', absolute_address, None descriptor = self._address_to_region_id[base_address] else: if target_region_id == 'global': # Just return the absolute address return 'global', absolute_address, None if target_region_id not in self._region_id_to_address: raise SimRegionMapError('Trying to relativize to a non-existent region "%s"' % target_region_id) descriptor = self._region_id_to_address[target_region_id] base_address = descriptor.base_address return descriptor.region_id, absolute_address - base_address, descriptor.related_function_address
class QLinearViewer(QWidget): def __init__(self, workspace, disasm_view, parent=None): super(QLinearViewer, self).__init__(parent) self.workspace = workspace self.disasm_view = disasm_view self.objects = [] # Objects that will be painted self.cfg = None self.cfb = None self._offset_to_addr = AVLTree() self._addr_to_offset = AVLTree() self._offset_to_object = AVLTree() self._offset = 0 self._paint_start_offset = 0 self._linear_view = None # type: QLinearGraphicsView self._disasms = {} self._init_widgets() # # Properties # @property def offset(self): return self._offset @offset.setter def offset(self, v): self._offset = v @property def paint_start_offset(self): return self._paint_start_offset @property def max_offset(self): # TODO: Cache it try: max_off, obj = self._offset_to_object.max_item() except (KeyError, ValueError): return 0 return max_off + obj.height # # Proxy properties # @property def selected_operands(self): return self._linear_view.selected_operands @property def selected_insns(self): return self._linear_view.selected_insns # # Public methods # def initialize(self): self._make_objects() def navigate_to_addr(self, addr): if not self._addr_to_offset: return try: _, floor_offset = self._addr_to_offset.floor_item(addr) except KeyError: _, floor_offset = floor_offset = self._addr_to_offset.min_item() self.navigate_to(floor_offset) def refresh(self): self._linear_view.refresh() def navigate_to(self, offset): self._linear_view.verticalScrollBar().setValue(int(offset)) self.prepare_objects(offset) self._linear_view.refresh() def prepare_objects(self, offset): if offset == self._offset: return try: start_offset = self._offset_to_object.floor_key(offset) except (KeyError, ValueError): try: start_offset = self._offset_to_object.min_key() except ValueError: # Tree is empty return # Update offset self._offset = offset self._paint_start_offset = start_offset self.objects = [] max_height = self.height() for off, obj in self._offset_to_object.iter_items( start_key=start_offset): self.objects.append(obj) if off - offset > max_height: break # # Private methods # def _init_widgets(self): self._linear_view = QLinearGraphicsView(self, self.disasm_view) layout = QHBoxLayout() layout.addWidget(self._linear_view) layout.setContentsMargins(0, 0, 0, 0) self.setLayout(layout) # Setup proxy methods self.update_label = self._linear_view.update_label self.select_instruction = self._linear_view.select_instruction self.unselect_instruction = self._linear_view.unselect_instruction self.unselect_all_instructions = self._linear_view.unselect_all_instructions self.select_operand = self._linear_view.select_operand self.unselect_operand = self._linear_view.unselect_operand self.unselect_all_operands = self._linear_view.unselect_all_operands self.show_selected = self._linear_view.show_selected self.show_instruction = self._linear_view.show_instruction def _make_objects(self): self._addr_to_offset.clear() self._offset_to_addr.clear() self._offset_to_object.clear() y = 0 self._linear_view._clear_insn_addr_block_mapping() for obj_addr, obj in self.cfb.floor_items(): if isinstance(obj, Block): func_addr = self.cfg.get_any_node( obj.addr).function_address # FIXME: Resiliency func = self.cfg.kb.functions[func_addr] # FIXME: Resiliency disasm = self._get_disasm(func) qobject = QBlock( self.workspace, func_addr, self.disasm_view, disasm, self.disasm_view.infodock, obj.addr, [obj], {}, mode='linear', ) for insn_addr in qobject.addr_to_insns.keys(): self._linear_view._add_insn_addr_block_mapping( insn_addr, qobject) elif isinstance(obj, Unknown): qobject = QUnknownBlock(self.workspace, obj_addr, obj.bytes) else: continue self._offset_to_object[y] = qobject if obj_addr not in self._addr_to_offset: self._addr_to_offset[obj_addr] = y y += qobject.height def _get_disasm(self, func): """ :param func: :return: """ if func.addr not in self._disasms: self._disasms[ func. addr] = self.workspace.instance.project.analyses.Disassembly( function=func) return self._disasms[func.addr]
class KeyedRegion(object): """ KeyedRegion keeps a mapping between stack offsets and all variables covering that offset. It assumes no variable in this region overlap with another variable in this region. Registers and function frames can all be viewed as a keyed region. """ def __init__(self, tree=None): self._storage = AVLTree() if tree is None else tree # type: AVLTree def __contains__(self, offset): """ Test if there is at least one varaible covering the given offset. :param offset: :return: """ try: base_offset, item = self._storage.floor_item(offset) #pylint:disable=unused-variable except KeyError: return False if item.includes(offset): return True return False def __len__(self): return len(self._storage) def __iter__(self): for _, item in self._storage.items(): yield item def __eq__(self, other): if set(self._storage.keys()) != set(other._storage.keys()): return False for k, v in self._storage.iter_items(): if v != other._storage[k]: return False return True def copy(self): if not self._storage: return KeyedRegion() kr = KeyedRegion() for key, ro in self._storage.iter_items(): kr._storage[key] = ro.copy() return kr def merge(self, other, make_phi_func=None): """ Merge another KeyedRegion into this KeyedRegion. :param KeyedRegion other: The other instance to merge with. :return: None """ # TODO: is the current solution not optimal enough? for _, item in other._storage.iter_items(): # type: RegionObject for loc_and_var in item.objects: self.__store(loc_and_var, overwrite=False, make_phi_func=make_phi_func) return self def dbg_repr(self): """ Get a debugging representation of this keyed region. :return: A string of debugging output. """ keys = self._storage.keys() offset_to_vars = {} for key in sorted(keys): ro = self._storage[key] variables = [obj.variable for obj in ro.objects] offset_to_vars[key] = variables s = [] for offset, variables in offset_to_vars.iteritems(): s.append("Offset %#x: %s" % (offset, variables)) return "\n".join(s) def add_variable(self, start, variable): """ Add a variable to this region at the given offset. :param int start: :param SimVariable variable: :return: None """ self._store(start, variable, overwrite=False) def set_variable(self, start, variable): """ Add a variable to this region at the given offset, and remove all other variables that are fully covered by this variable. :param int start: :param SimVariable variable: :return: None """ self._store(start, variable, overwrite=True) def get_base_addr(self, addr): """ Get the base offset (the key we are using to index variables covering the given offset) of a specific offset. :param int addr: :return: :rtype: int or None """ try: base_addr, item = self._storage.floor_item(addr) except KeyError: return None if item.includes(addr): return base_addr return None def get_variables_by_offset(self, start): """ Find variables covering the given region offset. :param int start: :return: A list of stack variables. :rtype: set """ try: base_addr = self._storage.floor_key(start) except KeyError: return [] item = self._storage[base_addr] # type: RegionObject if item.includes(start): return item.variables return [] # # Private methods # def _store(self, start, variable, overwrite=False): """ Store a variable into the storage. :param int start: The beginning address of the variable. :param variable: The variable to store. :param bool overwrite: Whether existing variables should be overwritten or not. :return: None """ loc_and_var = LocationAndVariable(start, variable) self.__store(loc_and_var, overwrite=overwrite) def __store(self, loc_and_var, overwrite=False, make_phi_func=None): """ Store a variable into the storage. :param LocationAndVariable loc_and_var: The descriptor describing start address and the variable. :param bool overwrite: Whether existing variables should be overwritten or not. :return: None """ start = loc_and_var.start variable = loc_and_var.variable variable_size = variable.size if variable.size is not None else 1 end = start + variable_size # region items in the middle overlapping_items = list(self._storage.item_slice(start, end)) # is there a region item that begins before the start and overlaps with this variable? try: floor_key, floor_item = self._storage.floor_item( start) # type: RegionObject if floor_item.includes(start): item = (floor_key, floor_item) if item not in overlapping_items: # insert it into the beginningq overlapping_items.insert(0, (floor_key, floor_item)) except KeyError: # no there isn't pass # scan through the entire list of region items, split existing regions and insert new regions as needed to_update = {start: RegionObject(start, variable_size, {loc_and_var})} last_end = start for _, item in overlapping_items: # type: RegionObject if item.start < start: # we need to break this item into two a, b = item.split(start) if overwrite: b.set_object(loc_and_var) else: self._add_object_or_make_phi(b, loc_and_var, make_phi_func=make_phi_func) to_update[a.start] = a to_update[b.start] = b last_end = b.end elif item.start > last_end: # there is a gap between the last item and the current item # fill in the gap new_item = RegionObject(last_end, item.start - last_end, {loc_and_var}) to_update[new_item.start] = new_item last_end = new_item.end elif item.end > end: # we need to split this item into two a, b = item.split(end) if overwrite: a.set_object(loc_and_var) else: self._add_object_or_make_phi(a, loc_and_var, make_phi_func=make_phi_func) to_update[a.start] = a to_update[b.start] = b last_end = b.end else: if overwrite: item.set_object(loc_and_var) else: self._add_object_or_make_phi(item, loc_and_var, make_phi_func=make_phi_func) to_update[loc_and_var.start] = item self._storage.update(to_update) def _is_overlapping(self, start, variable): if variable.size is not None: # make sure this variable does not overlap with any other variable end = start + variable.size try: prev_offset = self._storage.floor_key(end - 1) except KeyError: prev_offset = None if prev_offset is not None: if start <= prev_offset < end: return True prev_item = self._storage[prev_offset][0] prev_item_size = prev_item.size if prev_item.size is not None else 1 if start < prev_offset + prev_item_size < end: return True else: try: prev_offset = self._storage.floor_key(start) except KeyError: prev_offset = None if prev_offset is not None: prev_item = self._storage[prev_offset][0] prev_item_size = prev_item.size if prev_item.size is not None else 1 if prev_offset <= start < prev_offset + prev_item_size: return True return False def _add_object_or_make_phi(self, item, loc_and_var, make_phi_func=None): #pylint:disable=no-self-use if not make_phi_func or len({loc_and_var.variable} | item.variables) == 1: item.add_object(loc_and_var) else: # make a phi node item.set_object( LocationAndVariable( loc_and_var.start, make_phi_func(loc_and_var.variable, *item.variables)))
class CFBlanket(Analysis): """ A Control-Flow Blanket is a representation for storing all instructions, data entries, and bytes of a full program. """ def __init__(self, cfg=None): self._blanket = AVLTree() self._ffi = cffi.FFI() if cfg is not None: self._from_cfg(cfg) else: _l.debug( "CFG is not specified. Initialize CFBlanket from the knowledge base." ) for func in self.kb.functions.values(): self.add_function(func) def floor_item(self, addr): return self._blanket.floor_item(addr) def floor_items(self, addr=None): if addr is None: addr = self._blanket.min_key() else: try: addr = self.floor_addr(addr) except KeyError: try: addr = self.ceiling_addr(addr) except KeyError: # Nothing to yield raise StopIteration while True: try: item = self._blanket[addr] yield (addr, item) item_size = item.size if item.size > 0 else 1 # pylint: disable=no-member addr = self.ceiling_addr(addr + item_size) except KeyError: break def floor_addr(self, addr): return self._blanket.floor_key(addr) def ceiling_item(self, addr): return self._blanket.ceiling_item(addr) def ceiling_addr(self, addr): return self._blanket.ceiling_key(addr) def __getitem__(self, addr): return self._blanket[addr] def add_obj(self, addr, obj): """ :param addr: :param obj: :return: """ self._blanket[addr] = obj def add_function(self, func): """ Add a function and all blocks of this function to the blanket. :param angr.Function func: The function to add. :return: """ for block in func.blocks: self.add_obj(block.addr, block) def dbg_repr(self): """ The debugging representation of this CFBlanket. :return: The debugging representation of this CFBlanket. :rtype: str """ output = [] for obj in self.project.loader.all_objects: for section in obj.sections: if section.memsize == 0: continue min_addr, max_addr = section.min_addr, section.max_addr output.append("### Object %s" % repr(section)) output.append("### Range %#x-%#x" % (min_addr, max_addr)) pos = min_addr while pos < max_addr: try: addr, thing = self.floor_item(pos) output.append("%#x: %s" % (addr, repr(thing))) if thing.size == 0: pos += 1 else: pos += thing.size except KeyError: pos += 1 output.append("") return "\n".join(output) def _from_cfg(self, cfg): """ Initialize CFBlanket from a CFG instance. :param cfg: A CFG instance. :return: None """ # Let's first add all functions first for func in cfg.kb.functions.values(): self.add_function(func) self._mark_unknowns() def _mark_unknowns(self): """ Mark all unmapped regions. :return: None """ for obj in self.project.loader.all_objects: if isinstance(obj, cle.ELF): # sections? if obj.sections: for section in obj.sections: if not section.memsize or not section.vaddr: continue min_addr, max_addr = section.min_addr, section.max_addr self._mark_unknowns_core(min_addr, max_addr, obj=obj, section=section) elif obj.segments: for segment in obj.segments: if not segment.memsize: continue min_addr, max_addr = segment.min_addr, segment.max_addr self._mark_unknowns_core(min_addr, max_addr, obj=obj, segment=segment) else: # is it empty? _l.warning("Empty ELF object %s.", repr(obj)) elif isinstance(obj, cle.PE): if obj.sections: for section in obj.sections: if not section.memsize: continue min_addr, max_addr = section.min_addr, section.max_addr self._mark_unknowns_core(min_addr, max_addr, obj=obj, section=section) else: # is it empty? _l.warning("Empty PE object %s.", repr(obj)) else: min_addr, max_addr = obj.min_addr, obj.max_addr self._mark_unknowns_core(min_addr, max_addr, obj=obj) def _mark_unknowns_core(self, min_addr, max_addr, obj=None, segment=None, section=None): try: addr = self.floor_addr(min_addr) if addr < min_addr: raise KeyError except KeyError: # there is no other lower address try: next_addr = self.ceiling_addr(min_addr) if next_addr >= max_addr: raise KeyError except KeyError: next_addr = max_addr size = next_addr - min_addr if obj is None or isinstance(obj, cle.ExternObject): bytes_ = None else: try: _l.debug( "Loading bytes from object %s, section %s, segmeng %s, addresss %#x.", obj, section, segment, min_addr) bytes_ptr, _ = self.project.loader.memory.read_bytes_c( min_addr) bytes_ = self._ffi.unpack( self._ffi.cast('char*', bytes_ptr), size) # type: str except KeyError: # The address does not exist bytes_ = None self.add_obj( min_addr, Unknown(min_addr, size, bytes_=bytes_, object_=obj, segment=segment, section=section)) addr = min_addr while addr < max_addr: last_addr, last_item = self.floor_item(addr) if last_addr < min_addr: # impossible raise Exception('Impossible') if last_item.size == 0: # Make sure everything has a non-zero size last_item_size = 1 else: last_item_size = last_item.size end_addr = last_addr + last_item_size if end_addr < max_addr: try: next_addr = self.ceiling_addr(end_addr) except KeyError: next_addr = max_addr if next_addr > end_addr: # there is a gap size = next_addr - end_addr if obj is None or isinstance(obj, cle.ExternObject): bytes_ = None else: try: _l.debug( "Loading bytes from object %s, section %s, segmeng %s, addresss %#x.", obj, section, segment, next_addr) bytes_ptr, _ = self.project.loader.memory.read_bytes_c( next_addr) bytes_ = self._ffi.unpack( self._ffi.cast('char*', bytes_ptr), size) # type: str except KeyError: # The address does not exist bytes_ = None self.add_obj( end_addr, Unknown(end_addr, size, bytes_=bytes_, object_=obj, segment=segment, section=section)) addr = next_addr else: addr = max_addr
class KeyedRegion(object): """ KeyedRegion keeps a mapping between stack offsets and all variables covering that offset. It assumes no variable in this region overlap with another variable in this region. Registers and function frames can all be viewed as a keyed region. """ def __init__(self, tree=None): self._storage = AVLTree() if tree is None else tree # type: AVLTree def __contains__(self, offset): """ Test if there is at least one varaible covering the given offset. :param offset: :return: """ try: base_offset, item = self._storage.floor_item(offset) #pylint:disable=unused-variable except KeyError: return False if item.includes(offset): return True return False def __len__(self): return len(self._storage) def __iter__(self): for _, item in self._storage.items(): yield item def __eq__(self, other): if set(self._storage.keys()) != set(other._storage.keys()): return False for k, v in self._storage.iter_items(): if v != other._storage[k]: return False return True def copy(self): if not self._storage: return KeyedRegion() kr = KeyedRegion() for key, ro in self._storage.iter_items(): kr._storage[key] = ro.copy() return kr def merge(self, other, make_phi_func=None): """ Merge another KeyedRegion into this KeyedRegion. :param KeyedRegion other: The other instance to merge with. :return: None """ # TODO: is the current solution not optimal enough? for _, item in other._storage.iter_items(): # type: RegionObject for loc_and_var in item.objects: self.__store(loc_and_var, overwrite=False, make_phi_func=make_phi_func) return self def dbg_repr(self): """ Get a debugging representation of this keyed region. :return: A string of debugging output. """ keys = self._storage.keys() offset_to_vars = { } for key in sorted(keys): ro = self._storage[key] variables = [ obj.variable for obj in ro.objects ] offset_to_vars[key] = variables s = [ ] for offset, variables in offset_to_vars.iteritems(): s.append("Offset %#x: %s" % (offset, variables)) return "\n".join(s) def add_variable(self, start, variable): """ Add a variable to this region at the given offset. :param int start: :param SimVariable variable: :return: None """ self._store(start, variable, overwrite=False) def set_variable(self, start, variable): """ Add a variable to this region at the given offset, and remove all other variables that are fully covered by this variable. :param int start: :param SimVariable variable: :return: None """ self._store(start, variable, overwrite=True) def get_base_addr(self, addr): """ Get the base offset (the key we are using to index variables covering the given offset) of a specific offset. :param int addr: :return: :rtype: int or None """ try: base_addr, item = self._storage.floor_item(addr) except KeyError: return None if item.includes(addr): return base_addr return None def get_variables_by_offset(self, start): """ Find variables covering the given region offset. :param int start: :return: A list of stack variables. :rtype: set """ try: base_addr = self._storage.floor_key(start) except KeyError: return [ ] item = self._storage[base_addr] # type: RegionObject if item.includes(start): return item.variables return [ ] # # Private methods # def _store(self, start, variable, overwrite=False): """ Store a variable into the storage. :param int start: The beginning address of the variable. :param variable: The variable to store. :param bool overwrite: Whether existing variables should be overwritten or not. :return: None """ loc_and_var = LocationAndVariable(start, variable) self.__store(loc_and_var, overwrite=overwrite) def __store(self, loc_and_var, overwrite=False, make_phi_func=None): """ Store a variable into the storage. :param LocationAndVariable loc_and_var: The descriptor describing start address and the variable. :param bool overwrite: Whether existing variables should be overwritten or not. :return: None """ start = loc_and_var.start variable = loc_and_var.variable variable_size = variable.size if variable.size is not None else 1 end = start + variable_size # region items in the middle overlapping_items = list(self._storage.item_slice(start, end)) # is there a region item that begins before the start and overlaps with this variable? try: floor_key, floor_item = self._storage.floor_item(start) # type: RegionObject if floor_item.includes(start): item = (floor_key, floor_item) if item not in overlapping_items: # insert it into the beginningq overlapping_items.insert(0, (floor_key, floor_item)) except KeyError: # no there isn't pass # scan through the entire list of region items, split existing regions and insert new regions as needed to_update = { start: RegionObject(start, variable_size, { loc_and_var }) } last_end = start for _, item in overlapping_items: # type: RegionObject if item.start < start: # we need to break this item into two a, b = item.split(start) if overwrite: b.set_object(loc_and_var) else: self._add_object_or_make_phi(b, loc_and_var, make_phi_func=make_phi_func) to_update[a.start] = a to_update[b.start] = b last_end = b.end elif item.start > last_end: # there is a gap between the last item and the current item # fill in the gap new_item = RegionObject(last_end, item.start - last_end, { loc_and_var }) to_update[new_item.start] = new_item last_end = new_item.end elif item.end > end: # we need to split this item into two a, b = item.split(end) if overwrite: a.set_object(loc_and_var) else: self._add_object_or_make_phi(a, loc_and_var, make_phi_func=make_phi_func) to_update[a.start] = a to_update[b.start] = b last_end = b.end else: if overwrite: item.set_object(loc_and_var) else: self._add_object_or_make_phi(item, loc_and_var, make_phi_func=make_phi_func) to_update[loc_and_var.start] = item self._storage.update(to_update) def _is_overlapping(self, start, variable): if variable.size is not None: # make sure this variable does not overlap with any other variable end = start + variable.size try: prev_offset = self._storage.floor_key(end - 1) except KeyError: prev_offset = None if prev_offset is not None: if start <= prev_offset < end: return True prev_item = self._storage[prev_offset][0] prev_item_size = prev_item.size if prev_item.size is not None else 1 if start < prev_offset + prev_item_size < end: return True else: try: prev_offset = self._storage.floor_key(start) except KeyError: prev_offset = None if prev_offset is not None: prev_item = self._storage[prev_offset][0] prev_item_size = prev_item.size if prev_item.size is not None else 1 if prev_offset <= start < prev_offset + prev_item_size: return True return False def _add_object_or_make_phi(self, item, loc_and_var, make_phi_func=None): #pylint:disable=no-self-use if not make_phi_func or len({loc_and_var.variable} | item.variables) == 1: item.add_object(loc_and_var) else: # make a phi node item.set_object(LocationAndVariable(loc_and_var.start, make_phi_func(loc_and_var.variable, *item.variables) ) )