From 6f1d9b24da4a56e1480b7329887d67b56c502a1f Mon Sep 17 00:00:00 2001 From: Nikhil Parasaram Date: Thu, 18 Apr 2019 20:08:29 +0530 Subject: [PATCH] Support creation account --- mythril/analysis/symbolic.py | 1 + .../save_initial_world_state.py | 47 +++++++++++++++++++ .../laser/ethereum/plugins/plugin_factory.py | 8 ++++ mythril/laser/ethereum/svm.py | 42 ++++++++++------- 4 files changed, 80 insertions(+), 18 deletions(-) create mode 100644 mythril/laser/ethereum/plugins/implementations/save_initial_world_state.py diff --git a/mythril/analysis/symbolic.py b/mythril/analysis/symbolic.py index a7f54a77..969a9ae5 100644 --- a/mythril/analysis/symbolic.py +++ b/mythril/analysis/symbolic.py @@ -106,6 +106,7 @@ class SymExecWrapper: plugin_loader = LaserPluginLoader(self.laser) plugin_loader.load(PluginFactory.build_mutation_pruner_plugin()) plugin_loader.load(PluginFactory.build_instruction_coverage_plugin()) + plugin_loader.load(PluginFactory.build_set_initial_state_plugin()) self.laser.register_hooks( hook_type="pre", diff --git a/mythril/laser/ethereum/plugins/implementations/save_initial_world_state.py b/mythril/laser/ethereum/plugins/implementations/save_initial_world_state.py new file mode 100644 index 00000000..e3526d1f --- /dev/null +++ b/mythril/laser/ethereum/plugins/implementations/save_initial_world_state.py @@ -0,0 +1,47 @@ +from mythril.laser.ethereum.svm import LaserEVM +from mythril.laser.ethereum.plugins.plugin import LaserPlugin +from mythril.laser.ethereum.state.world_state import WorldState + +from typing import List + +import logging + +log = logging.getLogger(__name__) + + +class SaveInitialWorldState(LaserPlugin): + """SaveInitialWorldState + This plugin is used to save initial world state so it can be used for the output to display + + """ + + def __init__(self): + pass + + def initialize(self, symbolic_vm: LaserEVM): + """ + :param symbolic_vm: + :return: + """ + + @symbolic_vm.laser_hook("end_contract_creation") + def set_standard_initial_state(openstates: List[WorldState]): + """ + This function initializes the initial state to all the open states + :param openstates: + :return: + """ + accounts = openstates[0].accounts + initial_state = openstates[0].initial_state_account + initial_state[ + "accounts" + ] = {} # This variable persists for all world states. + for address, account in accounts.items(): + if address == "0x" + "0" * 40: + continue + initial_state["accounts"][address] = { + "nounce": account.nonce, + "balance": "", + "code": account.code.bytecode, + "storage": {}, + } diff --git a/mythril/laser/ethereum/plugins/plugin_factory.py b/mythril/laser/ethereum/plugins/plugin_factory.py index 1ce61b34..98a5011e 100644 --- a/mythril/laser/ethereum/plugins/plugin_factory.py +++ b/mythril/laser/ethereum/plugins/plugin_factory.py @@ -30,3 +30,11 @@ class PluginFactory: ) return InstructionCoveragePlugin() + + @staticmethod + def build_set_initial_state_plugin() -> LaserPlugin: + from mythril.laser.ethereum.plugins.implementations.save_initial_world_state import ( + SaveInitialWorldState, + ) + + return SaveInitialWorldState() diff --git a/mythril/laser/ethereum/svm.py b/mythril/laser/ethereum/svm.py index e52169e3..e535a71f 100644 --- a/mythril/laser/ethereum/svm.py +++ b/mythril/laser/ethereum/svm.py @@ -68,7 +68,6 @@ class LaserEVM: """ world_state = WorldState() world_state.accounts = accounts - world_state.initial_state_account = self.get_standard_initial_state(accounts) # this sets the initial world state self.world_state = world_state @@ -104,8 +103,20 @@ class LaserEVM: self._start_sym_exec_hooks = [] # type: List[Callable] self._stop_sym_exec_hooks = [] # type: List[Callable] + self._end_contract_creation_hooks = [] # type: List[Callable] + self.iprof = InstructionProfiler() if enable_iprof else None + self.laser_hooks_dict = { + "add_world_state": self._add_world_state_hooks, + "execute_state": self._execute_state_hooks, + "start_sym_exec": self._start_sym_exec_hooks, + "stop_sym_exec": self._stop_sym_exec_hooks, + "start_sym_trans": self._start_sym_trans_hooks, + "stop_sym_trans": self._stop_sym_trans_hooks, + "end_contract_creation": self._end_contract_creation_hooks, + } + log.info("LASER EVM initialized with dynamic loader: " + str(dynamic_loader)) @property @@ -116,17 +127,18 @@ class LaserEVM: """ return self.world_state.accounts - @staticmethod - def get_standard_initial_state(accounts: Dict[str, Account]) -> Dict: - template = {"accounts": {}} # type: Dict[str, Dict[str, Any]] + def set_standard_initial_state(self, accounts: Dict[str, Account]): + initial_state = self.world_state.initial_state_account + initial_state["accounts"] = {} # This variable persists for all world states. for address, account in accounts.items(): - template["accounts"][address] = { + if address == "0x" + "0" * 40: + continue + initial_state["accounts"][address] = { "nounce": account.nonce, "balance": "", "code": account.code.bytecode, "storage": {}, } - return template def sym_exec( self, main_address=None, creation_code=None, contract_name=None @@ -150,6 +162,7 @@ class LaserEVM: elif creation_code: log.info("Starting contract creation transaction") + created_account = execute_contract_creation( self, creation_code, contract_name ) @@ -163,6 +176,9 @@ class LaserEVM: "No contract was created during the execution of contract creation " "Increase the resources for creation execution (--max-depth or --create-timeout)" ) + else: + for hook in self._end_contract_creation_hooks: + hook(self.open_states) self._execute_transactions(created_account.address) @@ -496,18 +512,8 @@ class LaserEVM: def register_laser_hooks(self, hook_type: str, hook: Callable): """registers the hook with this Laser VM""" - if hook_type == "add_world_state": - self._add_world_state_hooks.append(hook) - elif hook_type == "execute_state": - self._execute_state_hooks.append(hook) - elif hook_type == "start_sym_exec": - self._start_sym_exec_hooks.append(hook) - elif hook_type == "stop_sym_exec": - self._stop_sym_exec_hooks.append(hook) - elif hook_type == "start_sym_trans": - self._start_sym_trans_hooks.append(hook) - elif hook_type == "stop_sym_trans": - self._stop_sym_trans_hooks.append(hook) + if hook_type in self.laser_hooks_dict: + self.laser_hooks_dict[hook_type].append(hook) else: raise ValueError( "Invalid hook type %s. Must be one of {add_world_state}", hook_type