From 6149e2911b872cfbd050367df4eaab75073567db Mon Sep 17 00:00:00 2001
From: Bernhard Mueller <bernhard.mueller11@gmail.com>
Date: Fri, 20 Oct 2017 09:45:40 +0700
Subject: [PATCH] Remove unneeded code from disassembler

---
 myth                                | 11 ++--
 mythril/disassembler/disassembly.py | 96 +----------------------------
 mythril/ether/evm.py                |  1 -
 tests/disassembler_test.py          |  5 +-
 4 files changed, 9 insertions(+), 104 deletions(-)

diff --git a/myth b/myth
index b0ea5a62..81c788b5 100755
--- a/myth
+++ b/myth
@@ -94,10 +94,13 @@ if (args.disassemble or args.graph or args.fire_lasers):
     else:
         exitWithError("No input bytecode. Please provide the code via -c BYTECODE or -a address")
 
-    try:
-        disassembly = Disassembly(encoded_bytecode)
-    except binascii.Error:
-        exitWithError("Disassembler: Invalid code string.")
+    if encoded_bytecode is not None:
+        logging.debug("Input bytecode: " + encoded_bytecode)
+
+        try:
+            disassembly = Disassembly(encoded_bytecode)
+        except binascii.Error:
+            exitWithError("Disassembler: Invalid code string.")
 
     if (args.disassemble):
 
diff --git a/mythril/disassembler/disassembly.py b/mythril/disassembler/disassembly.py
index e378cac5..2dc4b11d 100644
--- a/mythril/disassembler/disassembly.py
+++ b/mythril/disassembler/disassembly.py
@@ -3,32 +3,10 @@ import os
 import json
 
 
-class Block:
-
-    def __init__(self, id, code_index, funcname):
-        self.id = id
-        self.code_index = code_index
-        self.funcname = funcname
-        self.instruction_list = []
-
-    def update_length(self, num_instructions):
-        self.length = num_instructions
-        self.start_addr = self.instruction_list[0]['address']
-        self.end_addr = self.instruction_list[-1]['address']
-
-    def get_easm(self):
-
-        easm = str(self.instruction_list[0]['address']) + " " + self.funcname + "\n"
-        easm += asm.instruction_list_to_easm(self.instruction_list[1:])
-
-        return easm
-
-
 class Disassembly:
 
     def __init__(self, code):
         self.instruction_list = asm.disassemble(util.safe_decode(code))
-        self.blocks = []
         self.xrefs = []
         self.func_to_addr = {}
         self.addr_to_func = {}
@@ -59,80 +37,8 @@ class Disassembly:
             except:
                 continue
 
-        # Parse instructions into basic blocks
-
-        current_block = Block(0, 0, "PROLOGUE")
-
-        index = 0
-        blocklen = 0
-        blocknumber = 1
-
-        for instruction in self.instruction_list:
-
-            if (instruction['opcode'] == "JUMPDEST"):
-
-                try:
-                    func_name = "- FUNCTION " + self.addr_to_func[instruction['address']] + " -"
-                except KeyError:
-                    func_name = "- JUMPDEST_UNK -"
-                
-                current_block.update_length(blocklen)
-                self.blocks.append(current_block)
-                current_block = Block(blocknumber, index, func_name)
-                blocklen = 0
-                blocknumber += 1
-                
-            current_block.instruction_list.append(instruction)
-            blocklen += 1
-
-            index += 1
-        
-        # Add the last block
 
-        current_block.update_length(blocklen)
-        self.blocks.append(current_block)
-
-        # Resolve cross-references
-
-        for block in self.blocks:
-
-            jmp_indices = asm.find_opcode_sequence(["JUMP"], block.instruction_list)
-            jmp_indices += asm.find_opcode_sequence(["JUMPI"], block.instruction_list)
-
-            for i in jmp_indices:
-                try:
-                    dest_hex = block.instruction_list[i - 1]['argument']
-                    dest = int(dest_hex[2:], 16)
-                except:
-                    continue
-
-                j = 0
-
-                try:
-                    while(self.blocks[j].end_addr < dest):
-                        j += 1
-                except IndexError:
-                    continue
-
-
-                if not (block.id, self.blocks[j].id) in self.xrefs:
-                    self.xrefs.append((block.id, self.blocks[j].id))
-
-            # if the last instruction isn't an unconditional jump or halt, also add a reference to the following block
-
-            try:
-                if (block.id < len(self.blocks)) and (block.instruction_list[block.length - 1]['opcode'] not in ['JUMP', 'STOP', 'THROW', 'REVERT', 'INVALID']):
-                    if not (block.id, self.blocks[j].id) in self.xrefs:
-                        self.xrefs.append((block.id, block.id + 1))
-            except UnboundLocalError:
-                # quickfix
-                continue
 
     def get_easm(self):
 
-        easm = asm.instruction_list_to_easm(self.instruction_list[0:self.blocks[0].length])
-
-        for block in self.blocks[1:]:
-            easm += block.get_easm()
-
-        return easm
\ No newline at end of file
+        return asm.instruction_list_to_easm(self.instruction_list)
diff --git a/mythril/ether/evm.py b/mythril/ether/evm.py
index 42bbd737..f1e967a4 100644
--- a/mythril/ether/evm.py
+++ b/mythril/ether/evm.py
@@ -7,7 +7,6 @@ from io import StringIO
 import re
 
 
-
 def trace(code, calldata = ""):
 
 	logHandlers = ['eth.vm.op', 'eth.vm.op.stack', 'eth.vm.op.memory', 'eth.vm.op.storage']
diff --git a/tests/disassembler_test.py b/tests/disassembler_test.py
index abce8ede..172c162b 100644
--- a/tests/disassembler_test.py
+++ b/tests/disassembler_test.py
@@ -2,8 +2,6 @@ import unittest
 from mythril.disassembler.disassembly import Disassembly
 
 
-
-
 class ETHContractTestCase(unittest.TestCase):
 
     def setUp(self):
@@ -14,5 +12,4 @@ class DisassembyTestCase(ETHContractTestCase):
     def runTest(self):
 
         disassembly = Disassembly(self.code)
-
-        self.assertEqual(len(disassembly.blocks), 162, 'Disassembler error: Incorrect number of blocks generated)')
\ No newline at end of file
+        self.assertEqual(len(disassembly.instruction_list), 3537)