Merge pull request #1002 from ConsenSys/bugfix/source_list

Add the runtime bytecode to source list

mythril/analysis/symbolic.py

@@ -2,6 +2,7 @@
 purposes."""
 
 import copy
+from ethereum.utils import mk_contract_address
 from mythril.analysis.security import get_detection_module_hooks, get_detection_modules
 from mythril.laser.ethereum import svm
 from mythril.laser.ethereum.state.account import Account
@@ -11,11 +12,12 @@ from mythril.laser.ethereum.strategy.basic import (
     ReturnRandomNaivelyStrategy,
     ReturnWeightedRandomStrategy,
 )
+from mythril.laser.ethereum.transaction.symbolic import CREATOR_ADDRESS
+
 
 from mythril.laser.ethereum.plugins.plugin_factory import PluginFactory
 from mythril.laser.ethereum.plugins.plugin_loader import LaserPluginLoader
 
-
 from mythril.solidity.soliditycontract import EVMContract, SolidityContract
 from .ops import Call, SStore, VarType, get_variable
 
@@ -110,6 +112,11 @@ class SymExecWrapper:
             )
         else:
             self.laser.sym_exec(address)
+        created_address = "0x" + str(mk_contract_address(CREATOR_ADDRESS, 0).hex())
+        for key, value in self.laser.world_state.accounts.items():
+            if created_address == value.address:
+                contract.code = value.code.bytecode
+                break
 
         if not requires_statespace:
             return

mythril/disassembler/disassembly.py

@@ -28,11 +28,15 @@ class Disassembly(object):
         self.func_hashes = []  # type: List[str]
         self.function_name_to_address = {}  # type: Dict[str, int]
         self.address_to_function_name = {}  # type: Dict[int, str]
+        self.enable_online_lookup = enable_online_lookup
+        self.assign_bytecode(bytecode=code)
 
+    def assign_bytecode(self, bytecode):
+        self.bytecode = bytecode
         # open from default locations
         # control if you want to have online signature hash lookups
-        signatures = SignatureDB(enable_online_lookup=enable_online_lookup)
-
+        signatures = SignatureDB(enable_online_lookup=self.enable_online_lookup)
+        self.instruction_list = asm.disassemble(util.safe_decode(bytecode))
         # Need to take from PUSH1 to PUSH4 because solc seems to remove excess 0s at the beginning for optimizing
         jump_table_indices = asm.find_op_code_sequence(
             [("PUSH1", "PUSH2", "PUSH3", "PUSH4"), ("EQ",)], self.instruction_list

mythril/laser/ethereum/state/account.py

@@ -2,7 +2,7 @@
 
 This includes classes representing accounts and their storage.
 """
-
+from copy import deepcopy, copy
 from typing import Any, Dict, KeysView, Union
 
 from z3 import ExprRef
@@ -61,6 +61,13 @@ class Storage:
         """
         return self._storage.keys()
 
+    def __deepcopy__(self, memodict={}):
+        storage = Storage(
+            concrete=self.concrete, address=self.address, dynamic_loader=self.dynld
+        )
+        storage._storage = copy(self._storage)
+        return storage
+
 
 class Account:
     """Account class representing ethereum accounts."""
@@ -92,7 +99,6 @@ class Account:
         self.storage = Storage(
             concrete_storage, address=address, dynamic_loader=dynamic_loader
         )
-
         # Metadata
         self.address = address
         self.contract_name = contract_name
@@ -128,3 +134,14 @@ class Account:
             "balance": self.balance,
             "storage": self.storage,
         }
+
+    def __deepcopy__(self, memodict={}):
+        new_account = Account(
+            address=self.address,
+            code=self.code,
+            balance=self.balance,
+            contract_name=self.contract_name,
+        )
+        new_account.storage = deepcopy(self.storage)
+        new_account.code = self.code
+        return new_account

mythril/laser/ethereum/state/world_state.py

@@ -2,7 +2,7 @@
 from copy import copy
 from random import randint
 from typing import Dict, List, Iterator, Optional, TYPE_CHECKING
-
+from ethereum.utils import mk_contract_address
 from mythril.laser.ethereum.state.account import Account
 from mythril.laser.ethereum.state.annotation import StateAnnotation
 
@@ -50,7 +50,12 @@ class WorldState:
         return new_world_state
 
     def create_account(
-        self, balance=0, address=None, concrete_storage=False, dynamic_loader=None
+        self,
+        balance=0,
+        address=None,
+        concrete_storage=False,
+        dynamic_loader=None,
+        creator=None,
     ) -> Account:
         """Create non-contract account.
 
@@ -60,7 +65,7 @@ class WorldState:
         :param dynamic_loader: used for dynamically loading storage from the block chain
         :return: The new account
         """
-        address = address if address else self._generate_new_address()
+        address = address if address else self._generate_new_address(creator)
         new_account = Account(
             address,
             balance=balance,
@@ -111,11 +116,14 @@ class WorldState:
         """
         return filter(lambda x: isinstance(x, annotation_type), self.annotations)
 
-    def _generate_new_address(self) -> str:
+    def _generate_new_address(self, creator=None) -> str:
         """Generates a new address for the global state.
 
         :return:
         """
+        if creator:
+            # TODO: Use nounce
+            return "0x" + str(mk_contract_address(creator, 0).hex())
         while True:
             address = "0x" + "".join([str(hex(randint(0, 16)))[-1] for _ in range(40)])
             if address not in self.accounts.keys():

mythril/laser/ethereum/transaction/symbolic.py

@@ -73,7 +73,7 @@ def execute_contract_creation(
     del laser_evm.open_states[:]
 
     new_account = laser_evm.world_state.create_account(
-        0, concrete_storage=True, dynamic_loader=None
+        0, concrete_storage=True, dynamic_loader=None, creator=CREATOR_ADDRESS
     )
     if contract_name:
         new_account.contract_name = contract_name

mythril/laser/ethereum/transaction/transaction_models.py

@@ -188,7 +188,7 @@ class ContractCreationTransaction(BaseTransaction):
 
         contract_code = bytes.hex(array.array("B", return_data).tostring())
 
-        global_state.environment.active_account.code = Disassembly(contract_code)
+        global_state.environment.active_account.code.assign_bytecode(contract_code)
         self.return_data = global_state.environment.active_account.address
         assert global_state.environment.active_account.code.instruction_list != []
 

mythril/support/source_support.py

@@ -55,4 +55,9 @@ class Source:
         :param bytecode_hash: The contract hash
         :return: The index of the contract in the _source_hash list
         """
-        return self._source_hash.index(bytecode_hash)
+        # TODO: Add this part to exception logs
+        try:
+            return self._source_hash.index(bytecode_hash)
+        except ValueError:
+            self._source_hash.append(bytecode_hash)
+            return len(self._source_hash) - 1