<commit_before>import logging
from collections import defaultdict

import claripy

from .memory import SimMemory
from .symbolic_memory import SimSymbolicMemory

l = logging.getLogger("simuvex.plugins.abstract_memory")

class MemoryRegion(object):
    def __init__(self, id, state, init_memory=True, backer_dict=None):
        self._id = id
        self._state = state
        # This is a map from tuple (basicblock_key, stmt_id) to
        # AbstractLocation objects
        self._alocs = {}

        if init_memory:
            if backer_dict is None:
                self._memory = SimSymbolicMemory(memory_id=id)
            else:
                self._memory = SimSymbolicMemory(backer=backer_dict,
                                                 memory_id=id)

            self._memory.set_state(state)

    @property
    def id(self):
        return self._id

    @property
    def memory(self):
        return self._memory

    @property
    def state(self):
        return self._state

    def set_state(self, state):
        self._state = state
        self._memory.set_state(state)

    def copy(self):
        r = MemoryRegion(self._id, self.state, init_memory=False)
        r._memory = self.memory.copy()
        r._alocs = self._alocs.copy()
        return r

    def store(self, addr, data, bbl_addr, stmt_id):
        if bbl_addr is not None and stmt_id is not None:
            aloc_id = (bbl_addr, stmt_id)
            if aloc_id not in self._alocs:
                self._alocs[aloc_id] = self.state.se.AbstractLocation(bbl_addr,
                                                                      stmt_id,
                                                                      self.id,
                                                                      addr,
                                                                      len(data) / 8)

                return self.memory.store(addr, data)
            else:
                import ipdb; ipdb.set_trace()
                self._alocs[aloc_id].update(addr, len(data) / 8)
                return self.memory.store_with_merge(addr, data)

    def load(self, addr, size, bbl_addr, stmt_id):
        #if bbl_addr is not None and stmt_id is not None:

        return self.memory.load(addr, size)

class SimAbstractMemory(SimMemory):
    '''
    This is an implementation of the abstract store in paper [TODO].

    Some differences:
    # For stack variables, we map the absolute stack address to each region so
      that we can effectively trace stack accesses. When tracing into a new
      function, you should call set_stack_address_mapping() to create a new mapping.
      When exiting from a function, you should cancel the previous mapping by
      calling unset_stack_address_mapping().
      Currently this is only used for stack!
    '''
    def __init__(self, backer=None, memory_id="mem"):
        SimMemory.__init__(self)

        self._regions = {}
        self._stack_address_to_region = []
        self._stack_region_to_address = {}
        self._stack_size = None

        self._memory_id = memory_id

        if backer is not None:
            for region, backer_dict in backer.items():
                self._regions[region] = MemoryRegion(region, self.state,
                                               init_memory=True,
                                               backer_dict=backer_dict)

    def stack_id(self, function_address):
        return 'stack_0x%x' % function_address

    def set_stack_size(self, size):
        self._stack_size = size

    def set_stack_address_mapping(self, abs_addr, region_id):
        for address, region in self._stack_address_to_region:
            if address < abs_addr:
                self._stack_address_to_region.remove((address, region))
                del self._stack_region_to_address[region]

        self._stack_address_to_region.append((abs_addr, region_id))
        self._stack_region_to_address[region_id] = abs_addr

    def unset_stack_address_mapping(self, abs_addr, region_id):
        pos = self._stack_address_to_region.index((abs_addr, region_id))
        self._stack_address_to_region = self._stack_address_to_region[0 : pos]

        del self._stack_region_to_address[region_id]

    def _normalize_address(self, region, addr):
        '''
        If this is a stack address, we convert it to a correct region and address
        :param addr: Absolute address
        :return: a tuple of (region_id, normalized_address)
        '''
        stack_base = self._stack_address_to_region[0][0]

        if region.startswith('stack'):
            addr += self._stack_region_to_address[region]

            pos = 0
            for i in xrange(len(self._stack_address_to_region) - 1, 0, -1):
                if self._stack_address_to_region[i][0] > addr:
                    pos = i
                    break
            new_region = self._stack_address_to_region[pos][1]
            new_addr = addr - self._stack_address_to_region[pos][0]
            l.debug('%s 0x%08x is normalized to %s %08x, region base addr is 0x%08x', region, addr, new_region, new_addr, self._stack_address_to_region[pos][0])
            return (new_region, new_addr) # TODO: Is it OK to return a negative address?
        else:
            if addr == 0:
                import ipdb; ipdb.set_trace()
            l.debug("Got address %s 0x%x", region, addr)
            if addr < stack_base and \
                addr > stack_base - self._stack_size:
                return self._normalize_address('stack_initial', addr - stack_base)
            else:
                return (region, addr)

    def set_state(self, state):
        '''
        Overriding the SimStatePlugin.set_state() method
        :param state:
        :return:
        '''
        self.state = state
        for k, v in self._regions.items():
            v.set_state(state)

    # FIXME: symbolic_length is also a hack!
    def store(self, addr, data, size=None, key=None, condition=None, fallback=None, bbl_addr=None, stmt_id=None):
        if key is not None:
            raise DeprecationWarning('"key" is deprecated.')

        addr = addr._model
        assert type(addr) is claripy.vsa.ValueSet

        for region, addr_si in addr.items():
            # TODO: We only store to the min addr. Is this acceptable?
            normalized_region, normalized_addr = self._normalize_address(region, addr_si.min)
            self._store(normalized_addr, data, normalized_region, bbl_addr, stmt_id)

        # No constraints are generated...
        return []

    def _store(self, addr, data, key, bbl_addr, stmt_id):
        assert type(key) is str<commit_msg>Removed call to ipdb.set_trace()<commit_after>import logging
from collections import defaultdict

import claripy

from .memory import SimMemory
from .symbolic_memory import SimSymbolicMemory

l = logging.getLogger("simuvex.plugins.abstract_memory")

class MemoryRegion(object):
    def __init__(self, id, state, init_memory=True, backer_dict=None):
        self._id = id
        self._state = state
        # This is a map from tuple (basicblock_key, stmt_id) to
        # AbstractLocation objects
        self._alocs = {}

        if init_memory:
            if backer_dict is None:
                self._memory = SimSymbolicMemory(memory_id=id)
            else:
                self._memory = SimSymbolicMemory(backer=backer_dict,
                                                 memory_id=id)

            self._memory.set_state(state)

    @property
    def id(self):
        return self._id

    @property
    def memory(self):
        return self._memory

    @property
    def state(self):
        return self._state

    def set_state(self, state):
        self._state = state
        self._memory.set_state(state)

    def copy(self):
        r = MemoryRegion(self._id, self.state, init_memory=False)
        r._memory = self.memory.copy()
        r._alocs = self._alocs.copy()
        return r

    def store(self, addr, data, bbl_addr, stmt_id):
        if bbl_addr is not None and stmt_id is not None:
            aloc_id = (bbl_addr, stmt_id)
            if aloc_id not in self._alocs:
                self._alocs[aloc_id] = self.state.se.AbstractLocation(bbl_addr,
                                                                      stmt_id,
                                                                      self.id,
                                                                      addr,
                                                                      len(data) / 8)

                return self.memory.store(addr, data)
            else:
                self._alocs[aloc_id].update(addr, len(data) / 8)
                return self.memory.store_with_merge(addr, data)

    def load(self, addr, size, bbl_addr, stmt_id):
        #if bbl_addr is not None and stmt_id is not None:

        return self.memory.load(addr, size)

class SimAbstractMemory(SimMemory):
    '''
    This is an implementation of the abstract store in paper [TODO].

    Some differences:
    # For stack variables, we map the absolute stack address to each region so
      that we can effectively trace stack accesses. When tracing into a new
      function, you should call set_stack_address_mapping() to create a new mapping.
      When exiting from a function, you should cancel the previous mapping by
      calling unset_stack_address_mapping().
      Currently this is only used for stack!
    '''
    def __init__(self, backer=None, memory_id="mem"):
        SimMemory.__init__(self)

        self._regions = {}
        self._stack_address_to_region = []
        self._stack_region_to_address = {}
        self._stack_size = None

        self._memory_id = memory_id

        if backer is not None:
            for region, backer_dict in backer.items():
                self._regions[region] = MemoryRegion(region, self.state,
                                               init_memory=True,
                                               backer_dict=backer_dict)

    def stack_id(self, function_address):
        return 'stack_0x%x' % function_address

    def set_stack_size(self, size):
        self._stack_size = size

    def set_stack_address_mapping(self, abs_addr, region_id):
        for address, region in self._stack_address_to_region:
            if address < abs_addr:
                self._stack_address_to_region.remove((address, region))
                del self._stack_region_to_address[region]

        self._stack_address_to_region.append((abs_addr, region_id))
        self._stack_region_to_address[region_id] = abs_addr

    def unset_stack_address_mapping(self, abs_addr, region_id):
        pos = self._stack_address_to_region.index((abs_addr, region_id))
        self._stack_address_to_region = self._stack_address_to_region[0 : pos]

        del self._stack_region_to_address[region_id]

    def _normalize_address(self, region, addr):
        '''
        If this is a stack address, we convert it to a correct region and address
        :param addr: Absolute address
        :return: a tuple of (region_id, normalized_address)
        '''
        stack_base = self._stack_address_to_region[0][0]

        if region.startswith('stack'):
            addr += self._stack_region_to_address[region]

            pos = 0
            for i in xrange(len(self._stack_address_to_region) - 1, 0, -1):
                if self._stack_address_to_region[i][0] > addr:
                    pos = i
                    break
            new_region = self._stack_address_to_region[pos][1]
            new_addr = addr - self._stack_address_to_region[pos][0]
            l.debug('%s 0x%08x is normalized to %s %08x, region base addr is 0x%08x', region, addr, new_region, new_addr, self._stack_address_to_region[pos][0])
            return (new_region, new_addr) # TODO: Is it OK to return a negative address?
        else:
            l.debug("Got address %s 0x%x", region, addr)
            if addr < stack_base and \
                addr > stack_base - self._stack_size:
                return self._normalize_address(self._stack_address_to_region[0][1], addr - stack_base)
            else:
                return (region, addr)

    def set_state(self, state):
        '''
        Overriding the SimStatePlugin.set_state() method
        :param state:
        :return:
        '''
        self.state = state
        for k, v in self._regions.items():
            v.set_state(state)

    # FIXME: symbolic_length is also a hack!
    def store(self, addr, data, size=None, key=None, condition=None, fallback=None, bbl_addr=None, stmt_id=None):
        if key is not None:
            raise DeprecationWarning('"key" is deprecated.')

        addr = addr._model
        assert type(addr) is claripy.vsa.ValueSet

        for region, addr_si in addr.items():
            # TODO: We only store to the min addr. Is this acceptable?
            normalized_region, normalized_addr = self._normalize_address(region, addr_si.min)
            self._store(normalized_addr, data, normalized_region, bbl_addr, stmt_id)

        # No constraints are generated...
        return []

    def _store(self, addr, data, key, bbl_addr, stmt_id):
        assert type(key) is str
{'fraction_non_alphanumeric': 0.08167737060823574, 'fraction_numerical': 0.003702304495655459, 'mean_word_length': 1.9597495527728086, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 36, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2628873', 'n_tokens_mistral': 3787, 'n_tokens_neox': 3663, 'n_words': 981}