From 48948c4f7fca6df873824c9525151fa412bf6b40 Mon Sep 17 00:00:00 2001
From: Bernhard Mueller <bernhard.mueller11@gmail.com>
Date: Mon, 18 Sep 2017 14:52:49 +0700
Subject: [PATCH] Some refactoring

---
 code.easm    |  4 ++++
 e.asm        |  4 ++++
 ether/asm.py | 11 ++++-------
 lol          |  1 +
 mythril.py   | 18 ++++++++++++------
 5 files changed, 25 insertions(+), 13 deletions(-)
 create mode 100644 code.easm
 create mode 100644 e.asm
 create mode 100644 lol

diff --git a/code.easm b/code.easm
new file mode 100644
index 00000000..34a40a05
--- /dev/null
+++ b/code.easm
@@ -0,0 +1,4 @@
+PUSH2 0x4050
+PUSH4 0x60708090
+POP
+POP
diff --git a/e.asm b/e.asm
new file mode 100644
index 00000000..228867fd
--- /dev/null
+++ b/e.asm
@@ -0,0 +1,4 @@
+PUSH1 0x60
+PUSH1 0x40
+POP
+POP
diff --git a/ether/asm.py b/ether/asm.py
index f6176109..73892c26 100644
--- a/ether/asm.py
+++ b/ether/asm.py
@@ -1,7 +1,6 @@
 from ethereum import opcodes
 import codecs
 import re
-import binascii
 
 
 regex_PUSH = re.compile('^PUSH(\d*)$')
@@ -108,9 +107,7 @@ def resolve_functions(disassembly):
     return functions
 
 
-def disassemble(encoded_bytecode):
-
-    bytecode = safe_decode(encoded_bytecode)
+def disassemble(bytecode):
 
     disassembly = []
     i = 0
@@ -145,7 +142,7 @@ def disassemble(encoded_bytecode):
 
 def assemble(disassembly):
 
-    bytecode = ""
+    bytecode = b""
 
     for instruction in disassembly:
 
@@ -154,10 +151,10 @@ def assemble(disassembly):
         except RuntimeError:
             opcode = 0xbb
 
-        bytecode += binascii.hexlify(chr(opcode))
+        bytecode += chr(opcode)
 
         if 'argument' in instruction:
 
-            bytecode += instruction['argument']
+            bytecode += codecs.decode(instruction['argument'], 'hex_codec')
 
     return bytecode
diff --git a/lol b/lol
new file mode 100644
index 00000000..a11d5c87
--- /dev/null
+++ b/lol
@@ -0,0 +1 @@
+```@PP
\ No newline at end of file
diff --git a/mythril.py b/mythril.py
index ca32c3d5..2439b69f 100755
--- a/mythril.py
+++ b/mythril.py
@@ -6,6 +6,7 @@
 
 from ether import asm
 import sys
+import codecs
 import argparse
 import util
 
@@ -21,7 +22,7 @@ parser.add_argument('-d', '--disassemble',  action='store_true', help='disassemb
 parser.add_argument('-a', '--assemble', nargs=1, help='produce bytecode from easm input file', metavar='INPUT FILE')
 parser.add_argument('-c', '--code', nargs=1, help='bytecode string ("6060604052...")', metavar='BYTECODE')
 parser.add_argument('-t', '--transaction_hash', help='id of contract creation transaction')
-parser.add_argument('-o', '--outfile', help='file to write disassembly output to (e.g. "test.easm")')
+parser.add_argument('-o', '--outfile')
 parser.add_argument('--rpchost', nargs=1, help='RPC host')
 parser.add_argument('--rpcport', nargs=1, help='RPC port')
 
@@ -31,19 +32,19 @@ args = parser.parse_args()
 if (args.disassemble):
 
     if (args.code):
-        disassembly = asm.disassemble(args.code[0])
+        encoded_bytecode = args.code[0]
     elif (args.transaction_hash):
 
         try:
-            bytecode = util.bytecode_from_blockchain(args.transaction_hash)
+            encoded_bytecode = util.bytecode_from_blockchain(args.transaction_hash)
         except Exception as e:
             exitWithError("Exception loading bytecode via RPC: " + str(e.message))
 
-        disassembly = asm.disassemble(bytecode)
-
     else:
         exitWithError("Disassembler: Pass either the -c or -t flag to specify the input bytecode")
 
+    disassembly = asm.disassemble(util.safe_decode(encoded_bytecode))
+
     easm_text = asm.disassembly_to_easm(disassembly)
 
     if (args.outfile):
@@ -57,7 +58,12 @@ elif (args.assemble):
 
     disassembly = asm.easm_to_disassembly(easm)
 
-    print("0x" + asm.assemble(disassembly))
+    assembly = asm.assemble(disassembly)
+
+    if (args.outfile):
+        util.string_to_file(args.outfile, assembly)
+    else:
+        print("0x" + codecs.encode(assembly, "hex_codec"))
 
 else: