Repository: ConsenSys/mythril-classic
Pull request #603 (2018-10-22 21:59:03)

Refactor asm
The module is moved to disassembler, it seems like this is a more logical place for it to be.

Additionally some refactoring and code simplification

Todo:
- [x] Unit Tests


Patch:
diff --git a/mythril/analysis/symbolic.py b/mythril/analysis/symbolic.py
index 20790eb8..1108f4f7 100644
--- a/mythril/analysis/symbolic.py
+++ b/mythril/analysis/symbolic.py
@@ -4,8 +4,7 @@ from mythril.ether.soliditycontract import SolidityContract
 import copy
 import logging
 from .ops import get_variable, SStore, Call, VarType
-from mythril.laser.ethereum.strategy.basic import DepthFirstSearchStrategy, BreadthFirstSearchStrategy, \
-    ReturnRandomNaivelyStrategy, ReturnWeightedRandomStrategy
+from mythril.laser.ethereum.strategy.basic import DepthFirstSearchStrategy, BreadthFirstSearchStrategy
 
 
 class SymExecWrapper:
@@ -21,10 +20,6 @@ class SymExecWrapper:
             s_strategy = DepthFirstSearchStrategy
         elif strategy == 'bfs':
             s_strategy = BreadthFirstSearchStrategy
-        elif strategy == 'naive-random':
-            s_strategy = ReturnRandomNaivelyStrategy
-        elif strategy == 'weighted-random':
-            s_strategy = ReturnWeightedRandomStrategy
         else:
             raise ValueError("Invalid strategy argument supplied")
 
diff --git a/mythril/disassembler/asm.py b/mythril/disassembler/asm.py
new file mode 100644
index 00000000..762d1251
--- /dev/null
+++ b/mythril/disassembler/asm.py
@@ -0,0 +1,106 @@
+import re
+from collections import Generator
+
+from ethereum.opcodes import opcodes
+
+regex_PUSH = re.compile("^PUSH(\d*)$")
+
+# Additional mnemonic to catch failed assertions
+opcodes[254] = ["ASSERT_FAIL", 0, 0, 0]
+
+
+class EvmInstruction:
+    """ Model to hold the information of the disassembly """
+    def __init__(self, address, op_code, argument=None):
+        self.address = address
+        self.op_code = op_code
+        self.argument = argument
+
+    def to_dict(self) -> dict:
+        result = {"address": self.address, "opcode": self.op_code}
+        if self.argument:
+            result["argument"] = self.argument
+        return result
+
+
+def instruction_list_to_easm(instruction_list: dict) -> str:
+    result = ""
+
+    for instruction in instruction_list:
+        result += "{} {}".format(instruction["address"], instruction["opcode"])
+        if "argument" in instruction:
+            result += " " + instruction["argument"]
+        result += "\n"
+
+    return result
+
+
+def get_opcode_from_name(operation_name: str) -> int:
+    for op_code, value in opcodes.items():
+        if operation_name == value[0]:
+            return op_code
+    raise RuntimeError("Unknown opcode")
+
+
+def find_op_code_sequence(pattern: list, instruction_list: list) -> Generator:
+    """
+    Returns all indices in instruction_list that point to instruction sequences following a pattern
+    :param pattern: The pattern to look for.
+        Example:  [["PUSH1", "PUSH2"], ["EQ"]] where ["PUSH1", "EQ"] satisfies the pattern
+    :param instruction_list: List of instructions to look in
+    :return: Indices to the instruction sequences
+    """
+    for i in range(0, len(instruction_list) - len(pattern) + 1):
+        if is_sequence_match(pattern, instruction_list, i):
+            yield i
+
+
+def is_sequence_match(pattern: list, instruction_list: list, index: int) -> bool:
+    """
+    Checks if the instructions starting at index follow a pattern
+    :param pattern: List of lists describing a pattern.
+        Example:  [["PUSH1", "PUSH2"], ["EQ"]] where ["PUSH1", "EQ"] satisfies the pattern
+    :param instruction_list: List of instructions
+    :param index: Index to check for
+    :return: Pattern matched
+    """
+    for index, pattern_slot in enumerate(pattern, start=index):
+        try:
+            if not instruction_list[index]['opcode'] in pattern_slot:
+                return False
+        except IndexError:
+            return False
+    return True
+
+
+def disassemble(bytecode: str) -> list:
+    """Disassembles evm bytecode and returns a list of instructions"""
+    instruction_list = []
+    address = 0
+    length = len(bytecode)
+    if "bzzr" in str(bytecode[-43:]):
+        # ignore swarm hash
+        length -= 43
+
+    while address < length:
+        try:
+            op_code = opcodes[bytecode[address]]
+        except KeyError:
+            instruction_list.append(EvmInstruction(address, "INVALID"))
+            address += 1
+            continue
+
+        op_code_name = op_code[0]
+        current_instruction = EvmInstruction(address, op_code_name)
+
+        match = re.search(regex_PUSH, op_code_name)
+        if match:
+            argument_bytes = bytecode[address + 1: address + 1 + int(match.group(1))]
+            current_instruction.argument = "0x" + argument_bytes.hex()
+            address += int(match.group(1))
+
+        instruction_list.append(current_instruction)
+        address += 1
+
+    # We use a to_dict() here for compatibility reasons
+    return [element.to_dict() for element in instruction_list]
diff --git a/mythril/disassembler/disassembly.py b/mythril/disassembler/disassembly.py
index e28a4d73..d9529fe2 100644
--- a/mythril/disassembler/disassembly.py
+++ b/mythril/disassembler/disassembly.py
@@ -1,4 +1,5 @@
-from mythril.ether import asm, util
+from mythril.ether import util
+from mythril.disassembler import asm
 from mythril.support.signatures import SignatureDb
 import logging
 
@@ -21,7 +22,7 @@ class Disassembly(object):
         # Parse jump table & resolve function names
 
         # Need to take from PUSH1 to PUSH4 because solc seems to remove excess 0s at the beginning for optimizing
-        jmptable_indices = asm.find_opcode_sequence([("PUSH1", "PUSH2", "PUSH3", "PUSH4"), ("EQ",)],
+        jmptable_indices = asm.find_op_code_sequence([("PUSH1", "PUSH2", "PUSH3", "PUSH4"), ("EQ",)],
                                                     self.instruction_list)
 
         for i in jmptable_indices:
diff --git a/mythril/ether/asm.py b/mythril/ether/asm.py
deleted file mode 100644
index 985b2f07..00000000
--- a/mythril/ether/asm.py
+++ /dev/null
@@ -1,152 +0,0 @@
-import sys
-import re
-from ethereum.opcodes import opcodes
-from mythril.ether import util
-
-
-regex_PUSH = re.compile('^PUSH(\d*)$')
-
-# Additional mnemonic to catch failed assertions
-
-opcodes[254] = ['ASSERT_FAIL', 0, 0, 0]
-
-
-def instruction_list_to_easm(instruction_list):
-    easm = ""
-
-    for instruction in instruction_list:
-
-        easm += str(instruction['address']) + " " + instruction['opcode']
-
-        if 'argument' in instruction:
-            easm += " " + instruction['argument']
-
-        easm += "\n"
-
-    return easm
-
-
-def easm_to_instruction_list(easm):
-
-    regex_CODELINE = re.compile('^([A-Z0-9]+)(?:\s+([0-9a-fA-Fx]+))?$')
-
-    instruction_list = []
-
-    codelines = easm.split('\n')
-
-    for line in codelines:
-
-        m = re.search(regex_CODELINE, line)
-
-        if not m:
-            # Invalid code line
-            continue
-
-        instruction = {'opcode': m.group(1)}
-
-        if m.group(2):
-            instruction['argument'] = m.group(2)[2:]
-
-        instruction_list.append(instruction)
-
-    return instruction_list
-
-
-def get_opcode_from_name(name):
-
-    for opcode, value in opcodes.items():
-
-        if name == value[0]:
-
-            return opcode
-
-    raise RuntimeError("Unknown opcode")
-
-
-def find_opcode_sequence(pattern, instruction_list):
-    match_indexes = []
-
-    pattern_length = len(pattern)
-
-    for i in range(0, len(instruction_list) - pattern_length + 1):
-
-        if instruction_list[i]['opcode'] in pattern[0]:
-
-            matched = True
-
-            for j in range(1, len(pattern)):
-
-                if not (instruction_list[i + j]['opcode'] in pattern[j]):
-                    matched = False
-                    break
-
-            if matched:
-                match_indexes.append(i)
-
-    return match_indexes
-
-
-def disassemble(bytecode):
-
-    instruction_list = []
-    addr = 0
-
-    length = len(bytecode)
-
-    if "bzzr" in str(bytecode[-43:]):
-        # ignore swarm hash
-        length -= 43
-
-    while addr < length:
-
-        instruction = {'address': addr}
-
-        try:
-            if sys.version_info > (3, 0):
-                opcode = opcodes[bytecode[addr]]
-            else:
-                opcode = opcodes[ord(bytecode[addr])]
-
-        except KeyError:
-
-            # invalid opcode
-            instruction_list.append({'address': addr, 'opcode': "INVALID"})
-            addr += 1
-            continue
-
-        instruction['opcode'] = opcode[0]
-
-        m = re.search(regex_PUSH, opcode[0])
-
-        if m:
-            argument = bytecode[addr+1:addr+1+int(m.group(1))]
-            instruction['argument'] = "0x" + argument.hex()
-
-            addr += int(m.group(1))
-
-
-        instruction_list.append(instruction)
-
-        addr += 1
-
-    return instruction_list
-
-
-def assemble(instruction_list):
-
-    bytecode = b""
-
-    for instruction in instruction_list:
-
-        try:
-            opcode = get_opcode_from_name(instruction['opcode'])
-        except RuntimeError:
-            opcode = 0xbb
-
-        bytecode += opcode.to_bytes(1, byteorder='big')
-
-        if 'argument' in instruction:
-
-            bytecode += util.safe_decode(instruction['argument'])
-
-    return bytecode
diff --git a/mythril/interfaces/cli.py b/mythril/interfaces/cli.py
index ea55e29c..ae4a498d 100644
--- a/mythril/interfaces/cli.py
+++ b/mythril/interfaces/cli.py
@@ -70,11 +70,8 @@ def main():
     options = parser.add_argument_group('options')
     options.add_argument('-m', '--modules', help='Comma-separated list of security analysis modules', metavar='MODULES')
     options.add_argument('--max-depth', type=int, default=22, help='Maximum recursion depth for symbolic execution')
-
-    options.add_argument('--strategy', choices=['dfs', 'bfs', 'naive-random', 'weighted-random'],
-                         default='dfs', help='Symbolic execution strategy')
     options.add_argument('--max-transaction-count', type=int, default=3, help='Maximum number of transactions issued by laser')
-
+    options.add_argument('--strategy', choices=['dfs', 'bfs'], default='dfs', help='Symbolic execution strategy')
     options.add_argument('--execution-timeout', type=int, default=600, help="The amount of seconds to spend on symbolic execution")
     options.add_argument('--create-timeout', type=int, default=10, help="The amount of seconds to spend on "
                                                                         "the initial contract creation")
diff --git a/mythril/laser/ethereum/strategy/__init__.py b/mythril/laser/ethereum/strategy/__init__.py
index 3c208147..e69de29b 100644
--- a/mythril/laser/ethereum/strategy/__init__.py
+++ b/mythril/laser/ethereum/strategy/__init__.py
@@ -1,25 +0,0 @@
-from abc import ABC, abstractmethod
-
-
-class BasicSearchStrategy(ABC):
-    __slots__ = 'work_list', 'max_depth'
-
-    def __init__(self, work_list, max_depth):
-        self.work_list = work_list
-        self.max_depth = max_depth
-
-    def __iter__(self):
-        return self
-
-    @abstractmethod
-    def get_strategic_global_state(self):
-        raise NotImplementedError("Must be implemented by a subclass")
-
-    def __next__(self):
-        try:
-            global_state = self.get_strategic_global_state()
-            if global_state.mstate.depth >= self.max_depth:
-                return self.__next__()
-            return global_state
-        except IndexError:
-            raise StopIteration
diff --git a/mythril/laser/ethereum/strategy/basic.py b/mythril/laser/ethereum/strategy/basic.py
index 77c30b30..33dca443 100644
--- a/mythril/laser/ethereum/strategy/basic.py
+++ b/mythril/laser/ethereum/strategy/basic.py
@@ -1,67 +1,54 @@
 """
 This module implements basic symbolic execution search strategies
 """
-from random import randrange
-from . import BasicSearchStrategy
 
-try:
-    from random import choices
-except ImportError:
 
-    # This is for supporting python versions < 3.6
-    from itertools import accumulate
-    from random import random
-    from bisect import bisect
-
-    def choices(population, weights=None):
-        """
-        Returns a random element out of the population based on weight.
-        If the relative weights or cumulative weights are not specified,
-        the selections are made with equal probability.
-        """
-        if weights is None:
-            return [population[int(random() * len(population))]]
-        cum_weights = accumulate(weights)
-        return [population[bisect(cum_weights, random()*cum_weights[-1], 0, len(population)-1)]]
-
-
-class DepthFirstSearchStrategy(BasicSearchStrategy):
+class DepthFirstSearchStrategy:
     """
     Implements a depth first search strategy
     I.E. Follow one path to a leaf, and then continue to the next one
     """
+    def __init__(self, work_list, max_depth):
+        self.work_list = work_list
+        self.max_depth = max_depth
+
+    def __iter__(self):
+        return self
 
-    def get_strategic_global_state(self):
-        return self.work_list.pop()
+    def __next__(self):
+        """ Picks the next state to execute """
+        try:
+            # This strategies assumes that new states are appended at the end of the work_list
+            # By taking the last element we effectively pick the "newest" states, which amounts to dfs
+            global_state = self.work_list.pop()
+            if global_state.mstate.depth >= self.max_depth:
+                return self.__next__()
+            return global_state
+        except IndexError:
+            raise StopIteration()
 
 
-class BreadthFirstSearchStrategy(BasicSearchStrategy):
+class BreadthFirstSearchStrategy:
     """
     Implements a breadth first search strategy
     I.E. Execute all states of a "level" before continuing
     """
-
-    def get_strategic_global_state(self):
-        return self.work_list.pop(0)
-
-
-class ReturnRandomNaivelyStrategy(BasicSearchStrategy):
-    """
-    chooses a random state from the worklist with equal likelihood
-    """
-    def get_strategic_global_state(self):
-        if len(self.work_list) > 0:
-            return self.work_list.pop(randrange(len(self.work_list)))
-        else:
-            raise IndexError
-
-
-class ReturnWeightedRandomStrategy(BasicSearchStrategy):
-    """
-    chooses a random state from the worklist with likelihood based on inverse proportion to depth
-    """
-
-    def get_strategic_global_state(self):
-        probability_distribution = [1/(global_state.mstate.depth+1) for global_state in self.work_list]
-        return self.work_list.pop(choices(range(len(self.work_list)), probability_distribution)[0])
+    def __init__(self, work_list, max_depth):
+        self.work_list = work_list
+        self.max_depth = max_depth
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        """ Picks the next state to execute """
+        try:
+            # This strategies assumes that new states are appended at the end of the work_list
+            # By taking the first element we effectively pick the "oldest" states, which amounts to bfs
+            global_state = self.work_list.pop(0)
+            if global_state.mstate.depth >= self.max_depth:
+                return self.__next__()
+            return global_state
+        except IndexError:
+            raise StopIteration()
 


Test patch:
diff --git a/tests/disassembler/__init__.py b/tests/disassembler/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/disassembler/asm.py b/tests/disassembler/asm.py
new file mode 100644
index 00000000..960c27a9
--- /dev/null
+++ b/tests/disassembler/asm.py
@@ -0,0 +1,70 @@
+from mythril.disassembler.asm import *
+import pytest
+
+valid_names = [("PUSH1", 0x60), ("STOP", 0x0), ("RETURN", 0xf3)]
+
+
+@pytest.mark.parametrize("operation_name, hex_value", valid_names)
+def test_get_opcode(operation_name: str, hex_value: int):
+    # Act
+    return_value = get_opcode_from_name(operation_name)
+    # Assert
+    assert return_value == hex_value
+
+
+def test_get_unknown_opcode():
+    operation_name = "definitely unknown"
+
+    # Act
+    with pytest.raises(RuntimeError):
+        get_opcode_from_name(operation_name)
+
+
+sequence_match_test_data = [
+    # Normal no match
+    ((["PUSH1"], ["EQ"]), [{"opcode": "PUSH1"}, {"opcode": "PUSH3"}, {"opcode": "EQ"}], 1, False),
+    # Normal match
+    ((["PUSH1"], ["EQ"]), [{"opcode": "PUSH1"}, {"opcode": "PUSH1"}, {"opcode": "EQ"}], 1, True),
+    # Out of bounds pattern
+    ((["PUSH1"], ["EQ"]), [{"opcode": "PUSH1"}, {"opcode": "PUSH3"}, {"opcode": "EQ"}], 3, False),
+    ((["PUSH1"], ["EQ"]), [{"opcode": "PUSH1"}, {"opcode": "PUSH3"}, {"opcode": "EQ"}], 2, False),
+    # Double option match
+    ((["PUSH1", "PUSH3"], ["EQ"]), [{"opcode": "PUSH1"}, {"opcode": "PUSH1"}, {"opcode": "EQ"}], 1, True),
+    ((["PUSH1", "PUSH3"], ["EQ"]), [{"opcode": "PUSH1"}, {"opcode": "PUSH3"}, {"opcode": "EQ"}], 1, True),
+    # Double option no match
+    ((["PUSH1", "PUSH3"], ["EQ"]), [{"opcode": "PUSH1"}, {"opcode": "PUSH3"}, {"opcode": "EQ"}], 0, False),
+]
+
+
+@pytest.mark.parametrize("pattern, instruction_list, index,  expected_result", sequence_match_test_data)
+def test_is_sequence_match(pattern, instruction_list, index, expected_result):
+    # Act
+    return_value = is_sequence_match(pattern, instruction_list, index)
+    # Assert
+    assert return_value == expected_result
+
+find_sequence_match_test_data = [
+    # Normal no match
+    ((["PUSH1"], ["EQ"]), [{"opcode": "PUSH1"}, {"opcode": "PUSH3"}, {"opcode": "EQ"}], []),
+    # Normal match
+    ((["PUSH1"], ["EQ"]), [{"opcode": "PUSH1"}, {"opcode": "PUSH1"}, {"opcode": "EQ"}, {"opcode": "PUSH1"}, {"opcode": "EQ"}], [1, 3]),
+]
+
+
+@pytest.mark.parametrize("pattern, instruction_list, expected_result", find_sequence_match_test_data)
+def test_find_op_code_sequence(pattern, instruction_list, expected_result):
+    # Act
+    return_value = list(find_op_code_sequence(pattern, instruction_list))
+
+    # Assert
+    assert return_value == expected_result
+
+
+def test_disassemble():
+    # Act
+    instruction_list = disassemble(b"\x00\x16\x06")
+
+    # Assert
+    assert instruction_list[0]["opcode"] == "STOP"
+    assert instruction_list[1]["opcode"] == "AND"
+    assert instruction_list[2]["opcode"] == "MOD"