1import logging
2import claripy
3
4from angr.storage.memory_mixins import MemoryMixin
5
6l = logging.getLogger(__name__)
7
8
9class DataNormalizationMixin(MemoryMixin):
10    """
11    Normalizes the data field for a store and the fallback field for a load to be BVs.
12    """
13    def store(self, addr, data, size=None, **kwargs):
14        data_bv = self._convert_to_ast(data, size, self.state.arch.byte_width)
15
16        # zero extend if size is greater than len(data_e)
17        # TODO move this to the register resolver
18        #bit_width = size*self.state.arch.byte_width if isinstance(size, int) else self.state.arch.bits
19        #if size is not None and self.category == 'reg' and len(data_bv) < bit_width:
20        #    data_bv = data_bv.zero_extend(bit_width - len(data_bv))
21
22        if len(data_bv) % self.state.arch.byte_width != 0:
23            raise SimMemoryError("Attempting to store non-byte data to memory")
24
25        super().store(addr, data_bv, size=size, **kwargs)
26
27    def load(self, addr, size=None, fallback=None, **kwargs):
28        fallback_bv = self._convert_to_ast(fallback, size, self.state.arch.byte_width) if fallback is not None else None
29        return super().load(addr, size=size, fallback=fallback_bv, **kwargs)
30
31    def _convert_to_ast(self, thing, size, byte_width):
32        """
33        :param thing:       The thing to convert to an AST
34        :param size:        The size of the thing in bytes
35        :param byte_width:  The size of a byte in bits
36        """
37        if type(thing) is claripy.ast.BV:
38            return thing
39
40        if type(size) is int:
41            bits = size * byte_width
42        elif getattr(size, 'op', None) == 'BVV':
43            bits = size.args[0] * byte_width
44        else:
45            bits = None
46
47        if isinstance(thing, str):
48            l.warning("Encoding unicode string for memory as utf-8. Did you mean to use a bytestring?")
49            thing = thing.encode('utf-8')
50        if type(thing) is bytes:
51            return claripy.BVV(thing)
52        elif type(thing) is int:
53            if bits is None:
54                l.warning("Unknown size for memory data %#x. Default to arch.bits.", thing)
55                bits = self.state.arch.bits
56            return claripy.BVV(thing, bits)
57        elif type(thing) is float:
58            if bits == 32:
59                return claripy.FPV(thing, claripy.FSORT_FLOAT).raw_to_bv()
60            elif bits == 64:
61                return claripy.FPV(thing, claripy.FSORT_DOUBLE).raw_to_bv()
62            else:
63                raise TypeError("Passed float size which is not a float or a double", size)
64        else:
65            try:
66                raw_to_bv = thing.raw_to_bv
67            except AttributeError:
68                raise TypeError("Bad value passed to memory", thing) from None
69            else:
70                return raw_to_bv()
71
72from ...errors import SimMemoryError
73