# 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)