Refactor CLI

Author: defunctio

pyevmasm/__init__.py

@@ -1 +1 @@
-from .evmasm import EVMAsm
+from .evmasm import EVMAsm

evmasm renamed to pyevmasm/__main__.py

@@ -3,57 +3,53 @@
 import sys
 import binascii
 
-from pyevmasm import EVMAsm
+from .evmasm import EVMAsm
+
 
 def main():
     parser = argparse.ArgumentParser(description="pyevmasm the EVM assembler and disassembler")
-    parser.add_argument('-a','--assemble', action='store_true', help='Assemble EVM instructions to opcodes')
-    parser.add_argument('-d', '--disassemble', action='store_true', help='Disassemble EVM to opcodes')
-    parser.add_argument('-i', '--input', nargs='?', type=argparse.FileType('r'), default=sys.stdin, help='Input file, default=stdin')
-    parser.add_argument('-o', '--output', nargs='?', type=argparse.FileType('w'), default=sys.stdout, help='Output file, default=stdout')
-    parser.add_argument('-t', '--print-opcode-table', action='store_true', help='List supported EVM opcodes')
+    group_action = parser.add_mutually_exclusive_group(required=True)
+    group_action.add_argument('-a', '--assemble', action='store_true', help='Assemble EVM instructions to opcodes')
+    group_action.add_argument('-d', '--disassemble', action='store_true', help='Disassemble EVM to opcodes')
+    group_action.add_argument('-t', '--print-opcode-table', action='store_true', help='List supported EVM opcodes')
+    parser.add_argument('-i', '--input', nargs='?', type=argparse.FileType('r'), default=sys.stdin,
+                        help='Input file, default=stdin')
+    parser.add_argument('-o', '--output', nargs='?', type=argparse.FileType('w'), default=sys.stdout,
+                        help='Output file, default=stdout')
 
     args = parser.parse_args(sys.argv[1:])
 
     if args.print_opcode_table:
         table = EVMAsm._get_reverse_table()
         for mnemonic in table.keys():
             # This relies on the internal format 
-            (opcode, name, immediate_operand_size, pops, pushes, gas, description) =  table[mnemonic]
+            (opcode, name, immediate_operand_size, pops, pushes, gas, description) = table[mnemonic]
             print("%02x: %-16s %s" % (opcode, mnemonic, description))
-
         sys.exit(0)
 
-    if args.assemble and args.disassemble:
-        print("You cannot both assemble and disassemble at the same time.")
-        sys.exit(1)
-
-    if not args.assemble and not args.disassemble:
-        args.disassemble = True
-
     if args.assemble:
         asm = args.input.read().strip().rstrip()
         args.output.write(EVMAsm.assemble_hex(asm) + "\n")
 
     if args.disassemble:
         buf = args.input.read().strip().rstrip()
-        if buf[:3] == 'EVM': # binja prefix
+        if buf[:3] == 'EVM':  # binja prefix
             buf = buf[3:]
-        elif buf[:2] == '0x': # hex prefixed
+        elif buf[:2] == '0x':  # hex prefixed
             buf = binascii.unhexlify(buf[2:])
-        else: # detect all hex buffer
+        else:  # detect all hex buffer
             buf_set = set()
             for c in buf:
                 buf_set.add(c.lower())
 
             hex_set = set(list('0123456789abcdef'))
-            if buf_set <= hex_set: # subset
+            if buf_set <= hex_set:  # subset
                 buf = binascii.unhexlify(buf)
-        
+
         insns = list(EVMAsm.disassemble_all(buf))
         for i in insns:
-            args.output.write("%08x: %s\n" %(i.pc, str(i)))
+            args.output.write("%08x: %s\n" % (i.pc, str(i)))
+
 
 if __name__ == "__main__":
     main()
-

pyevmasm/evmasm.py

@@ -32,6 +32,7 @@ class EVMAsm(object):
             >>> EVMAsm.disassemble_hex('0x606040526002610100')
             'PUSH1 0x60\\nBLOCKHASH\\nMSTORE\\nPUSH1 0x2\\nPUSH2 0x100'
     '''
+
     class Instruction(object):
         def __init__(self, opcode, name, operand_size, pops, pushes, fee, description, operand=None, pc=0):
             '''
@@ -89,19 +90,20 @@ def __init__(self, opcode, name, operand_size, pops, pushes, fee, description, o
 
         def __eq__(self, other):
             ''' Instructions are equal if all features match '''
-            return self._opcode == other._opcode and\
-                self._name == other._name and\
-                self._operand == other._operand and\
-                self._operand_size == other._operand_size and\
-                self._pops == other._pops and\
-                self._pushes == other._pushes and\
-                self._fee == other._fee and\
-                self._pc == other._pc and\
-                self._description == other._description
+            return self._opcode == other._opcode and \
+                   self._name == other._name and \
+                   self._operand == other._operand and \
+                   self._operand_size == other._operand_size and \
+                   self._pops == other._pops and \
+                   self._pushes == other._pushes and \
+                   self._fee == other._fee and \
+                   self._pc == other._pc and \
+                   self._description == other._description
 
         def __repr__(self):
-            output = 'Instruction(0x%x, %r, %d, %d, %d, %d, %r, %r, %r)' % (self._opcode, self._name, self._operand_size,
-                                                                            self._pops, self._pushes, self._fee, self._description, self._operand, self._pc)
+            output = 'Instruction(0x%x, %r, %d, %d, %d, %d, %r, %r, %r)' % (
+            self._opcode, self._name, self._operand_size,
+            self._pops, self._pushes, self._fee, self._description, self._operand, self._pc)
             return output
 
         def __str__(self):
@@ -294,7 +296,8 @@ def uses_block_info(self):
         @property
         def is_arithmetic(self):
             ''' True if the instruction is an arithmetic operation '''
-            return self.semantics in ('ADD', 'MUL', 'SUB', 'DIV', 'SDIV', 'MOD', 'SMOD', 'ADDMOD', 'MULMOD', 'EXP', 'SIGNEXTEND')
+            return self.semantics in (
+            'ADD', 'MUL', 'SUB', 'DIV', 'SDIV', 'MOD', 'SMOD', 'ADDMOD', 'MULMOD', 'EXP', 'SIGNEXTEND')
 
     # from http://gavwood.com/paper.pdf
     _table = {  # opcode: (name, immediate_operand_size, pops, pushes, gas, description)
@@ -327,7 +330,8 @@ def is_arithmetic(self):
         0x31: ('BALANCE', 0, 1, 1, 20, 'Get balance of the given account.'),
         0x32: ('ORIGIN', 0, 0, 1, 2, 'Get execution origination address.'),
         0x33: ('CALLER', 0, 0, 1, 2, 'Get caller address.'),
-        0x34: ('CALLVALUE', 0, 0, 1, 2, 'Get deposited value by the instruction/transaction responsible for this execution.'),
+        0x34: (
+        'CALLVALUE', 0, 0, 1, 2, 'Get deposited value by the instruction/transaction responsible for this execution.'),
         0x35: ('CALLDATALOAD', 0, 1, 1, 3, 'Get input data of current environment.'),
         0x36: ('CALLDATASIZE', 0, 0, 1, 2, 'Get size of input data in current environment.'),
         0x37: ('CALLDATACOPY', 0, 3, 0, 3, 'Copy input data in current environment to memory.'),
@@ -336,7 +340,8 @@ def is_arithmetic(self):
         0x3a: ('GASPRICE', 0, 0, 1, 2, 'Get price of gas in current environment.'),
         0x3b: ('EXTCODESIZE', 0, 1, 1, 20, "Get size of an account's code."),
         0x3c: ('EXTCODECOPY', 0, 4, 0, 20, "Copy an account's code to memory."),
-        0x3d: ('RETURNDATASIZE', 0, 0, 1, 2, 'Get size of output data from the previous call from the current environment'),
+        0x3d: (
+        'RETURNDATASIZE', 0, 0, 1, 2, 'Get size of output data from the previous call from the current environment'),
         0x3e: ('RETURNDATACOPY', 0, 3, 0, 3, 'Copy output data from the previous call to memory'),
         0x40: ('BLOCKHASH', 0, 1, 1, 20, 'Get the hash of one of the 256 most recent complete blocks.'),
         0x41: ('COINBASE', 0, 0, 1, 2, "Get the block's beneficiary address."),
@@ -354,7 +359,8 @@ def is_arithmetic(self):
         0x57: ('JUMPI', 0, 2, 0, 10, 'Conditionally alter the program counter.'),
         0x58: ('GETPC', 0, 0, 1, 2, 'Get the value of the program counter prior to the increment.'),
         0x59: ('MSIZE', 0, 0, 1, 2, 'Get the size of active memory in bytes.'),
-        0x5a: ('GAS', 0, 0, 1, 2, 'Get the amount of available gas, including the corresponding reduction the amount of available gas.'),
+        0x5a: ('GAS', 0, 0, 1, 2,
+               'Get the amount of available gas, including the corresponding reduction the amount of available gas.'),
         0x5b: ('JUMPDEST', 0, 0, 0, 1, 'Mark a valid destination for jumps.'),
         0x60: ('PUSH', 1, 0, 1, 0, 'Place 1 byte item on stack.'),
         0x61: ('PUSH', 2, 0, 1, 0, 'Place 2-byte item on stack.'),
@@ -429,9 +435,11 @@ def is_arithmetic(self):
         0xf1: ('CALL', 0, 7, 1, 40, 'Message-call into an account.'),
         0xf2: ('CALLCODE', 0, 7, 1, 40, "Message-call into this account with alternative account's code."),
         0xf3: ('RETURN', 0, 2, 0, 0, 'Halt execution returning output data.'),
-        0xf4: ('DELEGATECALL', 0, 6, 1, 40, "Message-call into this account with an alternative account's code, but persisting into this account with an alternative account's code."),
+        0xf4: ('DELEGATECALL', 0, 6, 1, 40,
+               "Message-call into this account with an alternative account's code, but persisting into this account with an alternative account's code."),
         0xfa: ('STATICCALL', 0, 6, 1, 40, 'Static message-call into an account.'),
-        0xfd: ('REVERT', 0, 2, 0, 0, 'Stop execution and revert state changes, without consuming all provided gas and providing a reason.'),
+        0xfd: ('REVERT', 0, 2, 0, 0,
+               'Stop execution and revert state changes, without consuming all provided gas and providing a reason.'),
         0xfe: ('INVALID', 0, 0, 0, 0, 'Designated invalid instruction.'),
         0xff: ('SELFDESTRUCT', 0, 1, 0, 5000, 'Halt execution and register account for later deletion.')
     }
@@ -476,7 +484,8 @@ def assemble_one(assembler, pc=0):
                 assert len(assembler) == 1
                 operand = None
 
-            return EVMAsm.Instruction(opcode, name, operand_size, pops, pushes, gas, description, operand=operand, pc=pc)
+            return EVMAsm.Instruction(opcode, name, operand_size, pops, pushes, gas, description, operand=operand,
+                                      pc=pc)
         except BaseException:
             raise Exception("Something wrong at pc %d" % pc)
 
@@ -670,4 +679,4 @@ def assemble_hex(asmcode, pc=0):
                 ...
                 "0x6060604052600261010"
         '''
-        return '0x' + hexlify(EVMAsm.assemble(asmcode, pc=pc).encode()).decode()
+        return '0x' + hexlify(EVMAsm.assemble(asmcode, pc=pc).encode()).decode()

pyevmasm/util.py

@@ -33,4 +33,4 @@ def __repr__(self):
 
     def __get__(self, obj, objtype):
         '''Support instance methods.'''
-        return functools.partial(self.__call__, obj)
+        return functools.partial(self.__call__, obj)

setup.py

@@ -4,7 +4,6 @@
     name='pyevmasm',
     version='0.1.0',
     description='Ethereum Virtual Machine (EVM) assembler and disassembler',
-    scripts=['evmasm'],
     author='Trail of Bits',
     author_email='evmasm@trailofbits.com',
     url='https://github.com/trailofbits/pyevmasm',
@@ -13,10 +12,15 @@
     python_requires='>2.7',
     install_requires=[
         'future'
-        ],
+    ],
     extras_require={
         'dev': [
             'nose'
-            ]
-        }
+        ]
+    },
+    entry_points={
+        'console_scripts': [
+            'evmasm = pyevmasm.__main__:main'
+        ]
+    }
 )

tests/test_EVMAssembler.py

@@ -3,48 +3,46 @@
 from pyevmasm import EVMAsm
 
 
+# noinspection PyPep8Naming
 class EVMTest_Assembler(unittest.TestCase):
     _multiprocess_can_split_ = True
-    maxDiff=None 
+    maxDiff = None
 
     def test_ADD_1(self):
         instruction = EVMAsm.disassemble_one('\x60\x10')
         self.assertEqual(EVMAsm.Instruction(0x60, 'PUSH', 1, 0, 1, 0, 'Place 1 byte item on stack.', 16, 0),
                          instruction)
 
-
         instruction = EVMAsm.assemble_one('PUSH1 0x10')
         EVMAsm.Instruction(0x60, 'PUSH', 1, 0, 1, 0, 'Place 1 byte item on stack.', 16, 0)
-        
+
         instructions1 = EVMAsm.disassemble_all('\x30\x31')
         instructions2 = EVMAsm.assemble_all('ADDRESS\nBALANCE')
-        self.assertTrue( all(a == b for a,b in zip(instructions1, instructions2)))
+        self.assertTrue(all(a == b for a, b in zip(instructions1, instructions2)))
 
-        #High level simple assembler/disassembler
+        # High level simple assembler/disassembler
 
         bytecode = EVMAsm.assemble_hex(
-                                """PUSH1 0x60
-                                   BLOCKHASH
-                                   MSTORE
-                                   PUSH1 0x2
-                                   PUSH2 0x100
-                                """
-                             )
+            """PUSH1 0x60
+               BLOCKHASH
+               MSTORE
+               PUSH1 0x2
+               PUSH2 0x100
+            """
+        )
         self.assertEqual(bytecode, '0x606040526002610100')
 
-        asmcode =  EVMAsm.disassemble_hex('0x606040526002610100')
+        asmcode = EVMAsm.disassemble_hex('0x606040526002610100')
         self.assertEqual(asmcode, '''PUSH1 0x60\nBLOCKHASH\nMSTORE\nPUSH1 0x2\nPUSH2 0x100''')
 
-    def test_STOP(self):                                                                                                                    
-        insn = EVMAsm.disassemble_one('\x00')                                                                                               
-        self.assertTrue(str(insn) == 'STOP')                                                                                                
-                                                                                                                                            
-    def test_JUMPI(self):                                                                                                                   
-        insn = EVMAsm.disassemble_one('\x57')                                                                                               
-        self.assertTrue(str(insn) == 'JUMPI')                                                                                               
-        self.assertTrue(insn.is_branch)                                                                                                     
-
+    def test_STOP(self):
+        insn = EVMAsm.disassemble_one('\x00')
+        self.assertTrue(str(insn) == 'STOP')
 
+    def test_JUMPI(self):
+        insn = EVMAsm.disassemble_one('\x57')
+        self.assertTrue(str(insn) == 'JUMPI')
+        self.assertTrue(insn.is_branch)
 
 
 if __name__ == '__main__':