Concretise ITEStorage after each transaction

mythril/laser/ethereum/state/account.py

@@ -2,9 +2,10 @@
 
 This includes classes representing accounts and their storage.
 """
+from hashlib import sha3_256
 from copy import copy, deepcopy
 from typing import Any, Dict, Union, Tuple, cast, List
-
+from random import randint
 from mythril.laser.smt import (
     Array,
     K,
@@ -18,6 +19,7 @@ from mythril.laser.smt import (
 )
 from mythril.disassembler.disassembly import Disassembly
 from mythril.laser.smt import symbol_factory
+from z3.z3types import Z3Exception
 
 
 class StorageRegion:
@@ -181,7 +183,7 @@ class Storage:
                 and key.func_name == "keccak256"
                 and len(key.nested_functions) <= 1
             )
-            or key.symbolic is False
+            or key.potential_value is not None
         )
 
     def __getitem__(self, key: BitVec) -> BitVec:
@@ -209,6 +211,76 @@ class Storage:
         # TODO: Do something better here
         return str(self._printable_storage)
 
+    def concretize(self, model):
+        constraints = []
+        for key, value in self._ite_region.itelist:
+            if simplify(Extract(255, 0, key.input_)).symbolic or not isinstance(
+                key.input_, BitVecFunc
+            ):
+                continue
+            new_constraints, key = self._traverse_concretise(key, model)
+            constraints += new_constraints
+            self._array_region[key] = value
+        self._ite_region.itelist = []
+
+        return constraints
+
+    def calc_sha3(self, val):
+        try:
+            val = int(sha3_256(str(val.as_long()).encode("utf-8")).hexdigest(), 16)
+        except AttributeError:
+            val = int(
+                sha3_256(
+                    str(randint(0, 2 ** val.input_ - 1)).encode("utf-8")
+                ).hexdigest(),
+                16,
+            )
+        return symbol_factory.BitVecVal(val, 256)
+
+    def _find_value(self, symbol, model):
+        modify = symbol
+        if symbol.size() > 256:
+            index = simplify(Extract(255, 0, symbol))
+        else:
+            index = None
+        if index and not index.symbolic:
+            modify = Extract(511, 256, modify)
+        modify = model.eval(modify.raw)
+        try:
+            modify = modify.as_long()
+        except AttributeError:
+            modify = randint(0, 2 ** modify.size() - 1)
+        modify = symbol_factory.BitVecVal(modify, 256)
+        if index and not index.symbolic:
+            modify = Concat(modify, index)
+        return modify
+
+    def _traverse_concretise(self, key, model):
+        """
+        Breadth first Search
+        :param key:
+        :param model:
+        :return:
+        """
+        constraints = []
+        if not isinstance(key, BitVecFunc):
+            concrete_value = self._find_value(key, model)
+            constraints.append(concrete_value == key)
+            return constraints, concrete_value
+        for arg in key.concat_args:
+            new_const, val = self._traverse_concretise(arg, model)
+            constraints += new_const
+        if isinstance(key.input_, BitVec) or (
+            isinstance(key.input_, BitVecFunc) and key.input_.func_name == "sha3"
+        ):
+            new_const, _ = self._traverse_concretise(key.input_, model)
+            constraints += new_const
+
+        if key.func_name == "sha3":
+            key.potential_value = self.calc_sha3(key.input_)
+
+        return constraints, key
+
 
 class Account:
     """Account class representing ethereum accounts."""

mythril/laser/ethereum/svm.py

@@ -16,6 +16,8 @@ from mythril.laser.ethereum.state.world_state import WorldState
 from mythril.laser.ethereum.strategy.basic import DepthFirstSearchStrategy
 from abc import ABCMeta
 from mythril.laser.ethereum.time_handler import time_handler
+from mythril.laser.smt import Solver
+from z3 import unsat
 from mythril.laser.ethereum.transaction import (
     ContractCreationTransaction,
     TransactionEndSignal,
@@ -320,7 +322,7 @@ class LaserEVM:
             transaction, return_global_state = end_signal.global_state.transaction_stack[
                 -1
             ]
-
+            self.concretize_ite_storage(end_signal.global_state)
             if return_global_state is None:
                 if (
                     not isinstance(transaction, ContractCreationTransaction)
@@ -344,6 +346,17 @@ class LaserEVM:
 
         return new_global_states, op_code
 
+    def concretize_ite_storage(self, global_state):
+        solver = Solver()
+        for constraint in global_state.mstate.constraints:
+            solver.add(constraint)
+
+        if solver.check() == unsat:
+            return []
+        model = solver.model()
+        for account in global_state.world_state.accounts.values():
+            global_state.mstate.constraints += account.storage.concretize(model)
+
     def _end_message_call(
         self,
         return_global_state: GlobalState,
@@ -359,7 +372,6 @@ class LaserEVM:
         :param return_data:
         :return:
         """
-
         return_global_state.mstate.constraints += global_state.mstate.constraints
         # Resume execution of the transaction initializing instruction
         op_code = return_global_state.environment.code.instruction_list[

mythril/laser/smt/bitvec_helper.py

@@ -130,6 +130,9 @@ def Concat(*args: Union[BitVec, List[BitVec]]) -> BitVec:
         bvs = args[0]  # type: List[BitVec]
     else:
         bvs = cast(List[BitVec], args)
+    concat_list = []
+    for bv in bvs:
+        concat_list.append(bv)
 
     nraw = z3.Concat([a.raw for a in bvs])
     annotations = set()  # type: Annotations
@@ -147,6 +150,7 @@ def Concat(*args: Union[BitVec, List[BitVec]]) -> BitVec:
             func_name="Hybrid",
             input_=BitVec(z3.BitVec("", 256), annotations=annotations),
             nested_functions=nested_functions,
+            concat_args=concat_list,
         )
 
     return BitVec(nraw, annotations)
@@ -162,6 +166,25 @@ def Extract(high: int, low: int, bv: BitVec) -> BitVec:
     """
     raw = z3.Extract(high, low, bv.raw)
     if isinstance(bv, BitVecFunc):
+        count = 0
+        val = None
+        for small_bv in bv.concat_args[::-1]:
+            if low == count:
+                if low + small_bv.size() <= high:
+                    val = small_bv
+                else:
+                    val = Extract(low + small_bv.size() - 1, low, small_bv)
+            if high < count:
+                break
+            if low < count:
+                if low + small_bv.size() <= high:
+                    val = Concat(small_bv, val)
+                else:
+                    val = Extract(low + small_bv.size() - 1, low, small_bv)
+            count += small_bv.size()
+            val.simplify()
+        if val is not None:
+            return val
         input_string = ""
         # Is there a better value to set func_name and input to in this case?
         return BitVecFunc(

mythril/laser/smt/bitvecfunc.py

@@ -28,7 +28,7 @@ def _arithmetic_helper(
         return BitVecFunc(
             raw=raw,
             func_name="Hybrid",
-            input_=BitVec(z3.BitVec("", 256), annotations=union),
+            input_=BitVec(z3.BitVec("{} op {}".format(a, b), 256), annotations=union),
             nested_functions=a.nested_functions + b.nested_functions + [a, b],
         )
 
@@ -108,11 +108,14 @@ def _comparison_helper(
                 ),
             )
 
-    return And(
+    comparision = And(
         Bool(cast(z3.BoolRef, operation(a.raw, b.raw)), annotations=union),
         Bool(condition) if b.nested_functions else Bool(True),
         a.input_ == b.input_ if inputs_equal else a.input_ != b.input_,
     )
+    if a.potential_value is not None:
+        return Or(comparision, b == a.potential_value)
+    return comparision
 
 
 class BitVecFunc(BitVec):
@@ -125,6 +128,7 @@ class BitVecFunc(BitVec):
         input_: "BitVec" = None,
         annotations: Optional[Annotations] = None,
         nested_functions: Optional[List["BitVecFunc"]] = None,
+        concat_args: List = None,
     ):
         """
 
@@ -138,6 +142,8 @@ class BitVecFunc(BitVec):
         self.input_ = input_
         self.nested_functions = nested_functions or []
         self.nested_functions = list(dict.fromkeys(self.nested_functions))
+        self.potential_value = None
+        self.concat_args = concat_args or []
         if isinstance(input_, BitVecFunc):
             self.nested_functions.extend(input_.nested_functions)
         super().__init__(raw, annotations)