Source code for tritondse.memory

# built-in imports
import bisect
from typing import Optional, Union, Generator, List
from collections import namedtuple
import struct
from contextlib import contextmanager

# third-party imports
from triton import TritonContext

# local imports
from tritondse.types import Perm, Addr, ByteSize, Endian

MemMap = namedtuple('Map', "start size perm name")


class MapOverlapException(Exception):
    """
    Exception raised when trying to map a memory area where some of
    the addresses overlap with an already mapped area.
    """
    pass


class MemoryAccessViolation(Exception):
    """
    Exception triggered when accessing memory with
    the wrong permissions.
    """
    def __init__(self, addr: Addr, access: Perm, map_perm: Perm = None, memory_not_mapped: bool = False, perm_error: bool = False):
        """
        :param addr: address where the violation occurred
        :param access: type of access performed
        :param map_perm: permission of the memory page of `address`
        :param memory_not_mapped: whether the address was mapped or not
        :param perm_error: whether it is a permission error
        """
        super(MemoryAccessViolation, self).__init__()
        self.address: Addr = addr
        """ address where the violation occurred"""
        self._is_mem_unmapped = memory_not_mapped
        self._is_perm_error = perm_error
        self.access: Perm = access
        """Access type that was performed"""
        self.map_perm: Optional[Perm] = map_perm
        """Permissions of the memory map associated to the address"""

    def is_permission_error(self) -> bool:
        """True if the exception was caused by a permission issue"""
        return self._is_perm_error

    def is_memory_unmapped_error(self) -> bool:
        """
        Return true if the exception was raised due to a memory access
        to an area not mapped
        """
        return self._is_mem_unmapped

    def __str__(self) -> str:
        if self.is_permission_error():
            return f"(addr:{self.address:#08x}, access:{str(self.access)} on map:{str(self.map_perm)})"
        else:
            return f"({str(self.access)}: {self.address:#08x} unmapped)"

    def __repr__(self):
        return str(self)


STRUCT_MAP = {
    (True, 1): 'B',
    (False, 1): 'b',
    (True, 2): 'H',
    (False, 2): 'h',
    (True, 4): 'I',
    (False, 4): 'i',
    (True, 8): 'Q',
    (False, 8): 'q'
}

ENDIAN_MAP = {
    Endian.LITTLE: "<",
    Endian.BIG: ">"
}


[docs] class Memory(object): """ Memory representation of the current :py:class:`ProcessState` object. It wraps all interaction with Triton's memory context to provide high-level function. It adds a segmentation and memory permission model at the top of Triton. It also overrides __getitem__ and the slice mechanism to be able to read and write concrete memory values in a Pythonic manner. """ def __init__(self, ctx: TritonContext, endianness: Endian = Endian.LITTLE): """ :param ctx: TritonContext to interface with """ self.ctx: TritonContext = ctx """Underlying Triton context""" self._linear_map_addr = [] # List of [map_start, map_end, map_start, map_end ...] self._linear_map_map = [] # List of [MemMap, None, MemMap, None ...] self._segment_enabled = True self._endian = endianness self._endian_key = ENDIAN_MAP[self._endian] self._mem_cbs_enabled = True # self._maps = {} # Addr: -> Map
[docs] def set_endianness(self, en: Endian) -> None: """ Set the endianness of memory accesses. By default, endianness is little. :param en: Endian: Endianness to use. :return: None """ self._endian = en self._endian_key = ENDIAN_MAP[self._endian]
@property def _ptr_size(self) -> int: return self.ctx.getGprSize() @property def segmentation_enabled(self) -> bool: """ returns whether segmentation enforcing is enabled :return: True if segmentation is enabled """ return self._segment_enabled
[docs] def disable_segmentation(self) -> None: """ Turn-off segmentation enforcing. """ self._segment_enabled = False
[docs] def enable_segmentation(self) -> None: """ Turn-off segmentation enforcing. """ self._segment_enabled = True
[docs] def set_segmentation(self, enabled: bool) -> None: """ Set the segmentation enforcing with the given boolean. """ self._segment_enabled = enabled
[docs] @contextmanager def without_segmentation(self, disable_callbacks=False) -> Generator['Memory', None, None]: """ Context manager enabling manipulating temporarily the memory without considering the memory permissions. E.g: It enables writing data in a memory mapped in RX :param disable_callbacks: Whether to disable memory callbacks that could have been set :return: """ previous = self._segment_enabled self.disable_segmentation() cbs = self._mem_cbs_enabled self._mem_cbs_enabled = not disable_callbacks yield self self._mem_cbs_enabled = cbs self.set_segmentation(previous)
[docs] def callbacks_enabled(self) -> bool: """ Return whether memory callbacks are enabled. :return: True if callbacks are enabled """ return self._mem_cbs_enabled
[docs] def get_maps(self) -> Generator[MemMap, None, None]: """ Iterate all the memory maps defined, including all memory areas allocated on the heap. :return: generator of all :py:class:`MemMap` objects """ yield from (x for x in self._linear_map_map if x)
[docs] def map(self, start, size, perm: Perm = Perm.R | Perm.W | Perm.X, name="") -> MemMap: """ Map the given address and size in memory with the given permission. :raise MapOverlapException: In the case the map overlap an existing mapping :param start: address to map :param size: size to map :param perm: permission :param name: name to given to the memory region :return: MemMap freshly mapped """ def _map_idx(index): self._linear_map_addr.insert(index, start + size - 1) # end address is included self._linear_map_addr.insert(index, start) self._linear_map_map.insert(index, None) memmap = MemMap(start, size, perm, name) self._linear_map_map.insert(index, memmap) return memmap if not self._linear_map_addr: # Nothing mapped yet return _map_idx(0) idx = bisect.bisect_left(self._linear_map_addr, start) if idx == len(self._linear_map_addr): # It should be mapped at the end return _map_idx(idx) addr = self._linear_map_addr[idx] if (idx % 2) == 0: # We are on a start address if start < addr and start+size <= addr: # Can fit before return _map_idx(idx) else: # there is an overlap raise MapOverlapException(f"0x{start:08x}:{size} overlap with map: 0x{addr:08x} (even)") else: # We are on an end address prev = self._linear_map_addr[idx-1] raise MapOverlapException(f"0x{start:08x}:{size} overlap with map: 0x{prev:08x} (odd)")
[docs] def unmap(self, addr: Addr) -> None: """ Unmap the :py:class:`MemMap` object mapped at the address. The address can be within the map and not requires pointing at the head. :param addr: address to unmap :return: None """ def _unmap_idx(index): self._linear_map_addr.pop(index) # Pop the start self._linear_map_addr.pop(index) # Pop the end self._linear_map_map.pop(index) # Pop the object self._linear_map_map.pop(index) # Pop the None padding idx = bisect.bisect_left(self._linear_map_addr, addr) try: mapaddr = self._linear_map_addr[idx] if (idx % 2) == 0: # We are on a start address (meaning we should be exactly on map start other unmapped) if addr == mapaddr: # We are exactly on the map address _unmap_idx(idx) else: raise MemoryAccessViolation(addr, Perm(0), memory_not_mapped=True) else: # We are on an end address _unmap_idx(idx-1) except IndexError: raise MemoryAccessViolation(addr, Perm(0), memory_not_mapped=True)
[docs] def mprotect(self, addr: Addr, perm: Perm) -> None: """ Update the map at the given address with permissions provided in argument. :param addr: address of the map of which to change permission :param perm: permission to assign :return: None """ idx = bisect.bisect_left(self._linear_map_addr, addr) try: if (idx % 2) == 0: # We are on a start address (meaning we should be exactly on map start other unmapped) mmap = self._linear_map_map[idx] self._linear_map_map[idx] = MemMap(mmap.start, mmap.size, perm, mmap.name) # replace map with new perms else: # We are on an end address mmap = self._linear_map_map[idx-1] self._linear_map_map[idx-1] = MemMap(mmap.start, mmap.size, perm, mmap.name) # replace map with new perms except IndexError: raise MemoryAccessViolation(addr, Perm(0), memory_not_mapped=True)
def __setitem__(self, key: Addr, value: bytes) -> None: """ Assign the given value at the address given by the key. The value must be bytes but can be multiple bytes. Warning: You cannot use the slice API on this function. :param key: address to write to :param value: content to write :raise MemoryAccessViolation: in case of invalid access """ if isinstance(key, slice): raise TypeError("slice unsupported for __setitem__") else: self.write(key, value) def __getitem__(self, item: Union[Addr, slice]) -> bytes: """ Read the memory at the given address. If the key is an integer reads a single byte. If the key is a slice: read addr+size bytes in memory. :param item: address, or address:size to read :return: memory content :raise MemoryAccessViolation: if the access is invalid """ if isinstance(item, slice): return self.read(item.start, item.stop) elif isinstance(item, int): return self.read(item, 1)
[docs] def write(self, addr: Addr, data: bytes) -> None: """ Write the given `data` bytes at `addr` address. :param addr: address where to write :param data: data to write :return: None """ if self._segment_enabled: mmap = self._get_map(addr, len(data)) if mmap is None: raise MemoryAccessViolation(addr, Perm.W, memory_not_mapped=True) if Perm.W not in mmap.perm: raise MemoryAccessViolation(addr, Perm.W, map_perm=mmap.perm, perm_error=True) return self.ctx.setConcreteMemoryAreaValue(addr, data)
[docs] def read(self, addr: Addr, size: ByteSize) -> bytes: """ Read `size` bytes at `addr` address. :param addr: address to read :param size: size of content to read :return: bytes read """ if self._segment_enabled: mmap = self._get_map(addr, size) if mmap is None: raise MemoryAccessViolation(addr, Perm.R, memory_not_mapped=True) if Perm.R not in mmap.perm: raise MemoryAccessViolation(addr, Perm.R, map_perm=mmap.perm, perm_error=True) return self.ctx.getConcreteMemoryAreaValue(addr, size)
def _get_map(self, ptr: Addr, size: ByteSize) -> Optional[MemMap]: """ Internal function returning the MemMap object associated with any address. It returns None if part of the memory range falls out of a memory mapping. Complexity is O(log(n)) :param ptr: address in memory :param size: size of the memory :return: True if mapped """ idx = bisect.bisect_left(self._linear_map_addr, ptr) try: addr = self._linear_map_addr[idx] if (idx % 2) == 0: # We are on a start address (meaning we should be exactly on map start other unmapped) end = self._linear_map_addr[idx+1] return self._linear_map_map[idx] if (ptr == addr and ptr+size <= end+1) else None else: # We are on an end address start = self._linear_map_addr[idx-1] return self._linear_map_map[idx-1] if (start <= addr and ptr+size <= addr+1) else None # fit into the map except IndexError: return None # Either raised when linear_map is empty or the address is beyond everything that is mapped
[docs] def get_map(self, addr: Addr, size: ByteSize = 1) -> Optional[MemMap]: """ Find the MemMap associated with the given address and returns it if any. :param addr: Address of the map (or any map inside) :param size: size of bytes for which we want the map :return: MemMap if found """ return self._get_map(addr, size)
[docs] def find_map(self, name: str) -> Optional[List[MemMap]]: """ Find a map given its name. :param name: Map name :return: MemMap if found """ mmaps = [] for mmap in (x for x in self._linear_map_map if x): if mmap.name == name: mmaps.append(mmap) return mmaps
[docs] def map_from_name(self, name: str) -> MemMap: """ Return a map from its name. This function assumes the map is present. :raise AssertionError: If the map is not found :param name: Map name :return: MemMap """ for mmap in (x for x in self._linear_map_map if x): if mmap.name == name: return mmap assert False
[docs] def is_mapped(self, ptr: Addr, size: ByteSize = 1) -> bool: """ The function checks whether the memory is mapped or not. The implementation return False if the memory chunk overlap on two memory regions. Complexity is O(log(n)) :param ptr: address in memory :param size: size of the memory :return: True if mapped """ return self._get_map(ptr, size) is not None
[docs] def has_ever_been_written(self, ptr: Addr, size: ByteSize) -> bool: """ Returns whether the given range of addresses has previously been written or not. (Do not take in account the memory mapping). :param ptr: The pointer to check :type ptr: :py:obj:`tritondse.types.Addr` :param size: Size of the memory range to check :return: True if all addresses have been defined """ return self.ctx.isConcreteMemoryValueDefined(ptr, size)
[docs] def read_uint(self, addr: Addr, size: ByteSize = 4): """ Read in the process memory a **little-endian** integer of the ``size`` at ``addr``. :param addr: Address at which to read data :type addr: :py:obj:`tritondse.types.Addr` :param size: Number of bytes to read :type size: Union[str, :py:obj:`tritondse.types.ByteSize`] :return: Integer value read :raise struct.error: If value can't fit in `size` """ data = self.read(addr, size) return struct.unpack(self._endian_key+STRUCT_MAP[(True, size)], data)[0]
[docs] def read_sint(self, addr: Addr, size: ByteSize = 4): """ Read in the process memory a **little-endian** integer of the ``size`` at ``addr``. :param addr: Address at which to read data :type addr: :py:obj:`tritondse.types.Addr` :param size: Number of bytes to read :type size: Union[str, :py:obj:`tritondse.types.ByteSize`] :return: Integer value read :raise struct.error: If value can't fit in `size` """ data = self.read(addr, size) return struct.unpack(self._endian_key+STRUCT_MAP[(False, size)], data)[0]
[docs] def read_ptr(self, addr: Addr) -> int: """ Read in the process memory a little-endian integer of size :py:attr:`tritondse.ProcessState.ptr_size` :param addr: Address at which to read data :type addr: :py:obj:`tritondse.types.Addr` :return: Integer value read """ return self.read_uint(addr, self._ptr_size)
[docs] def read_char(self, addr: Addr) -> int: """ Read a char in memory (1-byte) following endianness. :param addr: address to read :return: char value as int """ return self.read_sint(addr, 1)
[docs] def read_uchar(self, addr: Addr) -> int: """ Read an unsigned char in memory (1-byte) following endianness. :param addr: address to read :return: unsigned char value as int """ return self.read_uint(addr, 1)
[docs] def read_int(self, addr: Addr) -> int: """ Read a signed integer in memory (4-byte) following endianness. :param addr: address to read :return: signed integer value as int """ return self.read_sint(addr, 4)
[docs] def read_word(self, addr: Addr) -> int: """ Read signed word in memory (2-byte) following endianness. :param addr: address to read :return: signed word value as int """ return self.read_uint(addr, 2)
[docs] def read_dword(self, addr: Addr) -> int: """ Read signed double word in memory (4-byte) following endianness. :param addr: address to read :return: dword value as int """ return self.read_uint(addr, 4)
[docs] def read_qword(self, addr: Addr) -> int: """ Read signed qword in memory (8-byte) following endianness. :param addr: address to read :return: qword value as int """ return self.read_uint(addr, 8)
[docs] def read_long(self, addr: Addr) -> int: """ Read 'C style' long in memory (4-byte) following endianness. :param addr: address to read :return: value as int """ return self.read_sint(addr, 4)
[docs] def read_ulong(self, addr: Addr) -> int: """ Read unsigned long in memory (4-byte) following endianness. :param addr: address to read :return: unsigned long value as int """ return self.read_uint(addr, 4)
[docs] def read_long_long(self, addr: Addr) -> int: """ Read long long in memory (8-byte) following endianness. :param addr: address to read :return: long long value as int """ return self.read_sint(addr, 8)
[docs] def read_ulong_long(self, addr: Addr) -> int: """ Read unsigned long long in memory (8-byte) following endianness. :param addr: address to read :return: unsigned long long value as int """ return self.read_uint(addr, 8)
[docs] def read_string(self, addr: Addr) -> str: """ Read a string in process memory at the given address .. warning:: The memory read is unbounded. Thus, the memory is iterated up until finding a 0x0. :returns: the string read in memory :rtype: str """ s = "" index = 0 while True: val = self.read_uint(addr+index, 1) if not val: return s s += chr(val) index += 1
[docs] def write_int(self, addr: Addr, value: int, size: ByteSize = 4): """ Write in the process memory the given integer value of the given size at a specific address. :param addr: Address at which to read data :param value: data to write represented as an integer :param size: Number of bytes to read :raise struct.error: If integer value cannot fit in `size` """ self.write(addr, struct.pack(self._endian_key+STRUCT_MAP[(value >= 0, size)], value))
[docs] def write_ptr(self, addr: Addr, value: int) -> None: """ Similar to :py:meth:`write_int` but the size is automatically adjusted to be ``ptr_size``. :param addr: address where to write data :type addr: :py:obj:`tritondse.types.Addr` :param value: pointer value to write :type value: int :raise struct.error: If integer value cannot fit in a pointer size """ self.write_int(addr, value, self._ptr_size)
[docs] def write_char(self, addr: Addr, value: int) -> None: """ Write the integer value as a single byte in memory. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a byte (>255) """ self.write_int(addr, value, 1)
[docs] def write_word(self, addr: Addr, value: int) -> None: """ Write the word (2-byte) in memory following endianness. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a word """ self.write_int(addr, value, 2)
[docs] def write_dword(self, addr: Addr, value: int) -> None: """ Write the word (4-byte) in memory following endianness. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a dword """ self.write_int(addr, value, 4)
[docs] def write_qword(self, addr: Addr, value: int) -> None: """ Write the qword (8-byte) in memory following endianness. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a qword """ self.write_int(addr, value, 8)
[docs] def write_long(self, addr: Addr, value: int) -> None: """ Write a "C style" long (4-byte) in memory following endianness. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a long """ return self.write_int(addr, value, 4)
[docs] def write_long_long(self, addr: Addr, value: int) -> None: """ Write the "C style" long long (8-byte) in memory following endianness. :param addr: address to write :param value: integer value :raise struct.error: If integer value do not fit in a long long """ return self.write_int(addr, value, 8)