Merge pull request #22 from cryptomental/feature/add-forks-support

pyevmasm, tests: Add support for multiple forks.

Author: disconnect3d

README.md

@@ -23,6 +23,9 @@ optional arguments:
                         Input file, default=stdin
   -o [OUTPUT], --output [OUTPUT]
                         Output file, default=stdout
+  -f FORK, --fork FORK  Fork, default: byzantium. Possible: [pre-byzantium,
+                        byzantium, constantinople]. Also an unsigned block
+                        number is accepted to select the fork.
 ```
 
 Disassembling the preamble of compiled contract:
@@ -44,7 +47,8 @@ $ echo -n "608060405260043610603f57600035" | evmasm -d
 ## Python API Examples
 
 ```
->>> from pyevmasm import instruction_table, disassemble_hex, disassemble_all, assemble_hex
+>>> from pyevmasm import instruction_tables, disassemble_hex, disassemble_all, assemble_hex
+>>> instruction_table = instruction_tables['byzantium']
 >>> instruction_table[20]
 Instruction(0x14, 'EQ', 0, 2, 1, 3, 'Equality comparision.', None, 0)
 >>> instruction_table['EQ']

pyevmasm/__init__.py

@@ -1,3 +1,4 @@
-from .evmasm import instruction_table, Instruction  # noqa: F401
+from .evmasm import instruction_tables, Instruction  # noqa: F401
+from .evmasm import block_to_fork, DEFAULT_FORK
 from .evmasm import assemble, assemble_all, assemble_hex, assemble_one
 from .evmasm import disassemble, disassemble_all, disassemble_hex, disassemble_one

pyevmasm/__main__.py

@@ -3,7 +3,7 @@
 import sys
 import binascii
 
-from .evmasm import assemble_hex, disassemble_all, instruction_table, assemble_all
+from .evmasm import assemble_hex, disassemble_all, instruction_tables, assemble_all, block_to_fork, DEFAULT_FORK
 
 
 def main():
@@ -18,9 +18,26 @@ def main():
                         help='Input file, default=stdin')
     parser.add_argument('-o', '--output', nargs='?', default=sys.stdout, type=argparse.FileType('w'),
                         help='Output file, default=stdout')
+    parser.add_argument('-f', '--fork', default=DEFAULT_FORK, type=str,
+                        help='Fork, default: byzantium. Possible: [pre-byzantium, byzantium, constantinople].'
+                             'Also an unsigned block number is accepted to select the fork.')
 
     args = parser.parse_args(sys.argv[1:])
 
+    accepted_forks = ['pre-byzantium', 'byzantium', 'constantinople']
+    arg_fork = args.fork.lower()
+    if arg_fork not in accepted_forks:
+        try:
+            block_number = abs(int(arg_fork))
+            fork = block_to_fork(block_number)
+        except ValueError:
+            sys.stderr.write('Wrong fork name or block number. '
+                             'Please provide an integer or one of %s.\n' % accepted_forks)
+            sys.exit(1)
+    else:
+        fork = arg_fork
+
+    instruction_table = instruction_tables[fork]
     if args.print_opcode_table:
         for instr in instruction_table:
             print('0x{:02x}: {:16s} {:s}'.format(instr.opcode, instr.name, instr.description))
@@ -32,13 +49,13 @@ def main():
         except KeyboardInterrupt:
             sys.exit(0)
         if args.binary_output:
-            for i in assemble_all(asm):
+            for i in assemble_all(asm, fork=fork):
                 if sys.version_info >= (3, 2):
                     args.output.buffer.write(i.bytes)
                 else:
                     args.output.write(i.bytes)
         else:
-            args.output.write(assemble_hex(asm) + "\n")
+            args.output.write(assemble_hex(asm, fork=fork) + "\n")
 
     if args.disassemble:
         if args.binary_input and sys.version_info >= (3, 2):
@@ -65,7 +82,7 @@ def main():
                 if buf_set <= hex_set:  # subset
                     buf = binascii.unhexlify(buf)
 
-        insns = list(disassemble_all(buf))
+        insns = list(disassemble_all(buf, fork=fork))
         for i in insns:
             args.output.write("%08x: %s\n" % (i.pc, str(i)))