bytecodeblock.py

from ceptions import InternalFunctionError
from bytecodes import *
from opcodes import INTERNAL_RETURN_OPCODE, exit_ops

import array, hashlib
import operator
from copy import deepcopy


class BytecodeBlock(object):
    def __init__(self, block_id):
        self.__block_id = block_id
        self.__items = list()

    def append(self, item):
        item.set_container_block_id(self.get_id())
        self.__items.append(item)

    def remove(self, index):
        del (self.__items[index])

    def insert(self, item, index):
        self.__items.insert(index, item)

    def get_id(self):
        return self.__block_id

    def get_block_size(self):
        return len(self.__items)

    def get_items(self):
        return self.__items

    def get_entry_address(self):
        return self.__items[0].get_address()

    def get_entry_bytecode(self):
        return self.__items[0]

    def get_exit_bytecode(self):
        return self.__items[-1]

    def is_exit_block(self):
        exit_bytecode = self.get_exit_bytecode()
        return exit_bytecode.opcode in exit_ops

    def is_jumpi_block(self):
        exit_bytecode = self.get_exit_bytecode()
        return exit_bytecode.opcode == "JUMPI"

    def is_jump_block(self):
        exit_bytecode = self.get_exit_bytecode()
        return exit_bytecode.opcode == "JUMP"

    def get_jumpi_address(self):
        if len(self.__items) < 2:
            return None
        jumpi_bytecode = self.__items[-1]
        push_bytecode = self.__items[-2]
        if jumpi_bytecode.opcode == "JUMPI" \
                and isinstance(push_bytecode, PushByteCode):
            return push_bytecode.get_value()

    def is_abort_block(self):
        exit_bytecode = self.__items[-1]
        if exit_bytecode.opcode in {"REVERT", "INVALID"}:
            return True
        if len(self.__items) < 2:
            return False
        push_bytecode = self.__items[-2]

        return exit_bytecode.opcode == "JUMP" \
               and isinstance(push_bytecode, PushByteCode) \
               and push_bytecode.get_value() in {2, 0}

    def get_function_signature(self):
        bytecodes = self.__items
        if len(self.__items) < 5: return -1
        bytecode_1 = bytecodes[-5]
        bytecode_2 = bytecodes[-4]
        bytecode_3 = bytecodes[-3]
        bytecode_4 = bytecodes[-2]
        bytecode_5 = bytecodes[-1]
        if bytecode_1.opcode == "DUP1" \
                and bytecode_2.opcode in {"PUSH4", "PUSH3"} \
                and bytecode_3.opcode == "EQ" \
                and isinstance(bytecode_4, PushByteCode) \
                and bytecode_5.opcode == "JUMPI":
            return bytecode_2.get_value()
        if bytecode_1.opcode in {"PUSH4", "PUSH3"} \
                and bytecode_2.opcode == "DUP2" \
                and bytecode_3.opcode == "EQ" \
                and isinstance(bytecode_4, PushByteCode) \
                and bytecode_5.opcode == "JUMPI":
            return bytecode_1.get_value()

        if len(self.__items) < 6: return -1
        bytecode_0 = bytecodes[-6]

        if bytecode_0.opcode == "DUP1" \
                and bytecode_1.opcode in {"PUSH4", "PUSH3"} \
                and bytecode_2.opcode == "EQ" \
                and bytecode_3.opcode == "ASSERT":
            return bytecode_1.get_value()
        if bytecode_0.opcode in {"PUSH4", "PUSH3"} \
                and bytecode_1.opcode == "DUP2" \
                and bytecode_2.opcode == "EQ" \
                and bytecode_3.opcode == "ASSERT":
            return bytecode_0.get_value()

        return -1

    def insert_assert(self):
        jumpi_bytecode = self.__items[-1]
        assert_bytecode = ByteCode("ASSERT", "", jumpi_bytecode.get_address(), -1)
        self.insert(assert_bytecode, -2)
        jumpi_bytecode.opcode = "JUMP"

    def insert_intcall(self, opcode, caller_end):
        new_block = self.make_copy()
        bytecodes = new_block.get_items()

        exit_bytecode = bytecodes[-1]
        if exit_bytecode.opcode == "JUMP":
            exit_bytecode.opcode = "POP"
        else:
            assert exit_bytecode.opcode != "JUMPI"
        intcall_bytecode = ByteCode(opcode, "gg", exit_bytecode.get_address(), 1)
        new_block.append(intcall_bytecode)
        push_bytecode = PushByteCode("PUSH1", "", exit_bytecode.get_address(), 2)
        address = caller_end.get_entry_address()
        push_bytecode.set_dependency(0, address)

        new_block.append(push_bytecode)
        jump_bytecode = ByteCode("JUMP", "", exit_bytecode.get_address(), 3)
        new_block.append(jump_bytecode)
        new_block.__block_id = self.get_id()
        # new_block.debug_block()

        return new_block

    def insert_intreturn(self):
        new_block = self.make_copy()
        bytecodes = new_block.get_items()
        bytecode_0 = bytecodes[-1]
        assert bytecode_0.opcode == "JUMP"

        new_block.__items[-1] = ByteCode("POP", "gg", bytecode_0.get_address())
        intreturn_bytecode = ByteCode(INTERNAL_RETURN_OPCODE, "gg", bytecode_0.get_address())
        new_block.append(intreturn_bytecode)
        new_block.__block_id = self.get_id()
        # new_block.debug_block()

        return new_block

    def merge(self, other):
        exit_bytecode = self.get_exit_bytecode()
        if exit_bytecode.opcode == "JUMP":
            exit_bytecode.opcode = "POP"
        for bytecode in other:
            self.append(bytecode)

    def is_empty(self):
        return len(self.__items) == 0

    def get_block_hash(self):
        raw_bytes = list()
        for bytecode in self:
            raw_byte = bytecode.raw_byte
            raw_bytes.append(raw_byte)
        raw_string = array.array('B', raw_bytes).tostring()
        m = hashlib.md5()
        m.update(raw_string)
        return int(m.hexdigest(), 16)

    def make_copy(self, block_id=-1):
        other = BytecodeBlock(block_id)
        other.__items = deepcopy(self.__items)
        return other

    def debug_block(self, depth=0):
        print("\nblock_%d" % self.get_id())
        for bytecode in self:
            print(bytecode)

    def __str__(self):
        return "block_%d" % self.get_id()

    def __iter__(self):
        for item in self.__items:
            yield item