manticore/native/state.py
import copy
import logging
from collections import namedtuple
from typing import Any, Callable, Dict, NamedTuple, Optional, Set, Tuple, Union
from .cpu.disasm import Instruction
from .memory import ConcretizeMemory, MemoryException
from .. import issymbolic
from ..core.state import StateBase, Concretize, TerminateState
from ..core.smtlib import Expression
from ..platforms import linux_syscalls
HookCallback = Callable[[StateBase], None]
logger = logging.getLogger(__name__)
class CheckpointData(NamedTuple):
pc: Any
last_pc: Any
class State(StateBase):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._hooks: Dict[Optional[int], Set[HookCallback]] = {}
self._after_hooks: Dict[Optional[int], Set[HookCallback]] = {}
self._sys_hooks: Dict[Optional[int], Set[HookCallback]] = {}
self._sys_after_hooks: Dict[Optional[int], Set[HookCallback]] = {}
def __getstate__(self) -> Dict[str, Any]:
state = super().__getstate__()
state["hooks"] = self._hooks
state["after_hooks"] = self._after_hooks
state["sys_hooks"] = self._sys_hooks
state["sys_after_hooks"] = self._sys_after_hooks
return state
def __setstate__(self, state: Dict[str, Any]) -> None:
super().__setstate__(state)
self._hooks = state["hooks"]
self._after_hooks = state["after_hooks"]
self._sys_hooks = state["sys_hooks"]
self._sys_after_hooks = state["sys_after_hooks"]
self._resub_hooks()
def __enter__(self) -> "State":
new_state = super().__enter__()
new_state._hooks = copy.copy(self._hooks)
new_state._after_hooks = copy.copy(self._after_hooks)
new_state._sys_hooks = copy.copy(self._sys_hooks)
new_state._sys_after_hooks = copy.copy(self._sys_after_hooks)
# Update constraint pointers in platform objects
from ..platforms.linux import SLinux
if isinstance(new_state.platform, SLinux):
from ..platforms.linux import SymbolicSocket
# Add constraints to symbolic sockets
for fd_entry in new_state.platform.fd_table.entries():
symb_socket_entry = fd_entry.fdlike
if isinstance(symb_socket_entry, SymbolicSocket):
symb_socket_entry._constraints = new_state.constraints
return new_state
def _get_hook_context(
self, after: bool = True, syscall: bool = False
) -> Tuple[Dict[Optional[int], Set[HookCallback]], str, Any]:
"""
Internal helper function to get hook context information.
:param after: Whether we want info pertaining to hooks after instruction executes or before
:param syscall: Catch a syscall invocation instead of instruction?
:return: Information for hooks after or before:
- set of hooks for specified after or before
- string of callback event
- State function that handles the callback
"""
if not syscall:
return (
(self._hooks, "will_execute_instruction", self._state_hook_callback)
if not after
else (self._after_hooks, "did_execute_instruction", self._state_after_hook_callback)
)
else:
return (
(self._sys_hooks, "will_invoke_syscall", self._state_sys_hook_callback)
if not after
else (
self._sys_after_hooks,
"did_invoke_syscall",
self._state_sys_after_hook_callback,
)
)
def remove_hook(
self,
pc_or_sys: Optional[Union[int, str]],
callback: HookCallback,
after: bool = False,
syscall: bool = False,
) -> bool:
"""
Remove a callback with the specified properties
:param pc_or_sys: Address of instruction, syscall number, or syscall name to remove hook from
:type pc_or_sys: int or None if `syscall` = False. int, str, or None if `syscall` = True
:param callback: The callback function that was at the address (or syscall)
:param after: Whether it was after instruction executed or not
:param syscall: Catch a syscall invocation instead of instruction?
:return: Whether it was removed
"""
if isinstance(pc_or_sys, str):
table = getattr(linux_syscalls, self._platform.current.machine)
for index, name in table.items():
if name == pc_or_sys:
pc_or_sys = index
break
if isinstance(pc_or_sys, str):
logger.warning(
f"{pc_or_sys} is not a valid syscall name in architecture {self._platform.current.machine}. "
"Please refer to manticore/platforms/linux_syscalls.py to find the correct name."
)
return False
hooks, when, _ = self._get_hook_context(after, syscall)
cbs = hooks.get(pc_or_sys, set())
if callback in cbs:
cbs.remove(callback)
else:
return False
if not len(hooks.get(pc_or_sys, set())):
del hooks[pc_or_sys]
return True
def add_hook(
self,
pc_or_sys: Optional[Union[int, str]],
callback: HookCallback,
after: bool = False,
syscall: bool = False,
) -> None:
"""
Add a callback to be invoked on executing a program counter (or syscall). Pass `None`
for `pc_or_sys` to invoke callback on every instruction (or syscall invocation).
`callback` should be a callable that takes one :class:`~manticore.native.state.State` argument.
:param pc_or_sys: Address of instruction to hook, syscall number, or syscall name
:type pc_or_sys: int or None if `syscall` = False. int, str, or None if `syscall` = True
:param callback: Hook function
:param after: Hook after PC (or after syscall) executes?
:param syscall: Catch a syscall invocation instead of instruction?
"""
if isinstance(pc_or_sys, str):
table = getattr(linux_syscalls, self._platform.current.machine)
for index, name in table.items():
if name == pc_or_sys:
pc_or_sys = index
break
if isinstance(pc_or_sys, str):
logger.warning(
f"{pc_or_sys} is not a valid syscall name in architecture {self._platform.current.machine}. "
"Please refer to manticore/platforms/linux_syscalls.py to find the correct name."
)
return
hooks, when, hook_callback = self._get_hook_context(after, syscall)
hooks.setdefault(pc_or_sys, set()).add(callback)
if hooks:
self.subscribe(when, hook_callback)
def _resub_hooks(self) -> None:
"""
Internal helper function to resubscribe hook callback events when the
state is active again.
"""
# TODO: check if the lists actually have hooks
_, when, hook_callback = self._get_hook_context(False, False)
self.subscribe(when, hook_callback)
_, when, hook_callback = self._get_hook_context(True, False)
self.subscribe(when, hook_callback)
_, when, hook_callback = self._get_hook_context(False, True)
self.subscribe(when, hook_callback)
_, when, hook_callback = self._get_hook_context(True, True)
self.subscribe(when, hook_callback)
def _state_hook_callback(self, pc: int, _instruction: Instruction) -> None:
"""
Invoke all registered State hooks before the instruction executes.
:param pc: Address where the hook should run
:param _instruction: Instruction at this PC
"""
# Prevent crash if removing hook(s) during a callback
tmp_hooks = copy.deepcopy(self._hooks)
# Invoke all pc-specific hooks
for cb in tmp_hooks.get(pc, []):
cb(self)
# Invoke all pc-agnostic hooks
for cb in tmp_hooks.get(None, []):
cb(self)
def _state_after_hook_callback(self, last_pc: int, _pc: int, _instruction: Instruction):
"""
Invoke all registered State hooks after the instruction executes.
:param last_pc: Address where the hook should run after instruction execution
:param _pc: Next address to execute
:param _instruction: Instruction at this last_pc
"""
# Prevent crash if removing hook(s) during a callback
tmp_hooks = copy.deepcopy(self._after_hooks)
# Invoke all pc-specific hooks
for cb in tmp_hooks.get(last_pc, []):
cb(self)
# Invoke all pc-agnostic hooks
for cb in tmp_hooks.get(None, []):
cb(self)
def _state_sys_hook_callback(self, syscall_num: int) -> None:
"""
Invoke all registered State hooks before the syscall executes.
:param syscall_num: index of the syscall about to be executed
"""
# Prevent crash if removing hook(s) during a callback
tmp_hooks = copy.deepcopy(self._sys_hooks)
# Invoke all syscall-specific hooks
for cb in tmp_hooks.get(syscall_num, []):
cb(self)
# Invoke all syscall-agnostic hooks
for cb in tmp_hooks.get(None, []):
cb(self)
def _state_sys_after_hook_callback(self, syscall_num: int):
"""
Invoke all registered State hooks after the syscall executes.
:param syscall_num: index of the syscall that was just executed
"""
# Prevent crash if removing hook(s) during a callback
tmp_hooks = copy.deepcopy(self._sys_after_hooks)
# Invoke all syscall-specific hooks
for cb in tmp_hooks.get(syscall_num, []):
cb(self)
# Invoke all syscall-agnostic hooks
for cb in tmp_hooks.get(None, []):
cb(self)
@property
def cpu(self):
"""
Current cpu state
"""
return self._platform.current
@property
def mem(self):
"""
Current virtual memory mappings
"""
return self._platform.current.memory
def _rollback(self, checkpoint_data: CheckpointData) -> None:
"""
Rollback state to previous values in checkpoint_data
"""
# Keep in this form to make sure we don't miss restoring any newly added
# data. Make sure the order is correct
self.cpu.PC, self.cpu._last_pc = checkpoint_data
def execute(self):
"""
Perform a single step on the current state
"""
super().execute()
from .cpu.abstractcpu import (
ConcretizeRegister,
) # must be here, otherwise we get circular imports
checkpoint_data = CheckpointData(pc=self.cpu.PC, last_pc=self.cpu._last_pc)
try:
result = self._platform.execute()
# Instead of State importing SymbolicRegisterException and SymbolicMemoryException
# from cpu/memory shouldn't we import Concretize from linux, cpu, memory ??
# We are forcing State to have abstractcpu
except ConcretizeRegister as exc:
# Need to define local variable to use in closure
e = exc
expression = self.cpu.read_register(e.reg_name)
def setstate(state: State, value):
state.cpu.write_register(e.reg_name, value)
self._rollback(checkpoint_data)
raise Concretize(str(e), expression=expression, setstate=setstate, policy=e.policy)
except ConcretizeMemory as exc:
# Need to define local variable to use in closure
e = exc
expression = self.cpu.read_int(e.address, e.size)
def setstate(state: State, value):
state.cpu.write_int(e.address, value, e.size)
self._rollback(checkpoint_data)
raise Concretize(str(e), expression=expression, setstate=setstate, policy=e.policy)
except Concretize as e:
self._rollback(checkpoint_data)
raise e
except MemoryException as e:
raise TerminateState(str(e), testcase=True)
# Remove when code gets stable?
assert self.platform.constraints is self.constraints
return result
def invoke_model(self, model):
"""
Invokes a `model`. Modelling can be used to override a function in the target program with a custom
implementation.
For more information on modelling see docs/models.rst
A `model` is a callable whose first argument is a `manticore.native.State` instance.
If the following arguments correspond to the arguments of the C function
being modeled. If the `model` models a variadic function, the following argument
is a generator object, which can be used to access function arguments dynamically.
The `model` callable should simply return the value that should be returned by the
native function being modeled.f
:param model: callable, model to invoke
"""
self._platform.invoke_model(model, prefix_args=(self,))
def _update_state_descriptor(self, descriptor, *args, **kwargs):
"""
Called on execution_intermittent to update the descriptor for this state.
This one should apply any native-specific information to the descriptor. Right now, that's just the PC and _last_pc
:param descriptor: StateDescriptor for this state
"""
super()._update_state_descriptor(descriptor, *args, **kwargs)
descriptor.pc = self.cpu.PC
descriptor.last_pc = self.cpu._last_pc