diff --git a/greed/factory.py b/greed/factory.py index 460898b..b8bc226 100644 --- a/greed/factory.py +++ b/greed/factory.py @@ -1,14 +1,22 @@ +import typing import logging from greed.block import Block from greed.sim_manager import SimulationManager from greed.state import SymbolicEVMState +if typing.TYPE_CHECKING: + from greed.project import Project + from greed.TAC import TAC_Statement + from greed.function import TAC_Function + log = logging.getLogger(__name__) class Factory: + project: 'Project' + def __init__(self, project: "Project"): self.project = project diff --git a/greed/sim_manager.py b/greed/sim_manager.py index 06a3cbb..ce8e1be 100644 --- a/greed/sim_manager.py +++ b/greed/sim_manager.py @@ -1,11 +1,14 @@ import logging import os import sys -from typing import Callable +from typing import Callable, List, Optional, TYPE_CHECKING from greed import options from greed.state import SymbolicEVMState +if TYPE_CHECKING: + from greed.exploration_techniques import ExplorationTechnique + log = logging.getLogger(__name__) @@ -35,35 +38,35 @@ def set_error(self, s: str): self.error += [s] @property - def states(self): + def states(self) -> List[SymbolicEVMState]: """ :return: All the states """ return sum(self.stashes.values(), []) @property - def active(self): + def active(self) -> List[SymbolicEVMState]: """ :return: Active stash """ return self.stashes['active'] @property - def deadended(self): + def deadended(self) -> List[SymbolicEVMState]: """ :return: Deadended stash """ return self.stashes['deadended'] @property - def found(self): + def found(self) -> List[SymbolicEVMState]: """ :return: Found stash """ return self.stashes['found'] @property - def one_active(self): + def one_active(self) -> Optional[SymbolicEVMState]: """ :return: First element of the active stash, or None if the stash is empty """ @@ -73,7 +76,7 @@ def one_active(self): return None @property - def one_deadended(self): + def one_deadended(self) -> Optional[SymbolicEVMState]: """ :return: First element of the deadended stash, or None if the stash is empty """ @@ -83,7 +86,7 @@ def one_deadended(self): return None @property - def one_found(self): + def one_found(self) -> Optional[SymbolicEVMState]: """ :return: First element of the found stash, or None if the stash is empty """ @@ -146,7 +149,7 @@ def step(self, find: Callable[[SymbolicEVMState], bool] = lambda s: False, for s in self.stashes['pruned'] + self.stashes['unsat'] + self.stashes['errored']: s.solver.dispose_context() - def single_step_state(self, state: SymbolicEVMState): + def single_step_state(self, state: SymbolicEVMState) -> List[SymbolicEVMState]: log.debug(f"Stepping {state}") log.debug(state.curr_stmt) diff --git a/greed/state.py b/greed/state.py index 544d62b..144d849 100644 --- a/greed/state.py +++ b/greed/state.py @@ -1,5 +1,6 @@ import datetime import logging +import typing from greed import options as opt from greed.memory import LambdaMemory, PartialConcreteStorage @@ -11,9 +12,28 @@ log = logging.getLogger(__name__) +if typing.TYPE_CHECKING: + from greed.project import Project + class SymbolicEVMState: uuid_generator = UUIDGenerator() + xid: int + project: "Project" + code: bytes + uuid: int + active_plugins: typing.Dict[str, SimStatePlugin] + _pc: str + trace: typing.List[str] + memory: LambdaMemory + options: typing.Dict[str, typing.Any] + registers: typing.Dict[str, typing.Any] + + # default plugins + solver: SimStateSolver + globals: SimStateGlobals + inspect: SimStateInspect + def __init__(self, xid, project, partial_init=False, init_ctx=None, options=None, max_calldatasize=None, partial_concrete_storage=False): self.xid = xid @@ -158,7 +178,7 @@ def set_init_ctx(self, init_ctx=None): self.ctx["CALLVALUE"] = BVV(init_ctx["CALLVALUE"], 256) @property - def pc(self): + def pc(self) -> str: return self._pc @property @@ -170,7 +190,7 @@ def constraints(self): return self.solver.constraints @pc.setter - def pc(self, value): + def pc(self, value: str): self._pc = value def set_next_pc(self): diff --git a/greed/state_plugins/solver.py b/greed/state_plugins/solver.py index 6a5a263..bc105d5 100644 --- a/greed/state_plugins/solver.py +++ b/greed/state_plugins/solver.py @@ -3,11 +3,14 @@ from greed import options from greed.solver.shortcuts import * from greed.state_plugins import SimStatePlugin +from greed.solver.solver import Solver log = logging.getLogger(__name__) class SimStateSolver(SimStatePlugin): + _solver: Solver + def __init__(self, partial_init=False): super(SimStateSolver, self).__init__() @@ -41,14 +44,14 @@ def _add_assertions(self, assertions): # Adding the constraints to the backend self._solver.add_assertions(assertions) - def push(self): + def push(self) -> int: self._curr_frame_level += 1 self._path_constraints[self._curr_frame_level] = set() self._memory_constraints[self._curr_frame_level] = set() self._solver.push() return self._curr_frame_level - def pop(self): + def pop(self) -> int: if self._curr_frame_level == 0: log.fatal("Cannot pop() level 0 frame") return self._curr_frame_level @@ -184,7 +187,7 @@ def eval_memory_at(self, array, offset, length, raw=False): else: return f"{int_value:0{bv_unsigned_value(length)*2}x}" - def copy(self): + def copy(self) -> 'SimStateSolver': new_solver = SimStateSolver(partial_init=True) new_solver._solver = self._solver.copy()