tracereader.py

from opcodes import *
from ceptions import TraceError, TraceException, ValidationError, InputError
from contextops import ContextOps, FreeOps, READ_POISON, EffectOps

from graphbuilder import FALLBACK_SIGNATURE

import sys, json


def load_trace(line):
    if len(line) > 1073741824:
        raise InputError("JSON too large")
    try:
        info = json.loads(line)
    except ValueError:
        raise InputError("JSON failed to parse")

    if 'error' in info:
        raise InputError("JSON missing transaction")
    return info


class TraceReader:
    def __init__(self, line):
        info = load_trace(line)
        self.code = info['code']
        self.__parse_signature(info)
        self.error = info['result']['failed']

        self.free_ops = FreeOps()
        self.memory_ops = ContextOps()
        self.storage_ops = ContextOps()

        t_ = info['result']['structLogs']
        self.trace = [i for i in t_ if i['depth'] == 1]

    def __parse_signature(self, info):
        signature = info['input_data'][:10]
        if signature == "0x":
            self.signature = FALLBACK_SIGNATURE
        else:
            self.signature = int(signature, 16)

    def parse_trace(self):
        self.__parse_trace()

    def __parse_trace(self):
        for index, item in enumerate(self.trace):
            opcode = item['op'].encode('utf-8')

            try:
                inputs, output = self.get_pair(opcode, index)
            except TraceException:
                self.trace = self.trace[:index]
                self.error = True
                break

            if opcode in free_ops:
                self.free_ops.add_mapping(inputs, output)
            if opcode in mem_read_ops:  # mem reads/writes are not mutually exclusive
                self.memory_ops.add_mapping(inputs, output)
            if opcode in mem_write_ops:  # there will be one in the context, and another as separator
                self.memory_ops.push_context(inputs)
            if opcode == "SLOAD":
                self.storage_ops.add_mapping(inputs, output)
            elif opcode == "SSTORE":
                self.storage_ops.push_context(inputs)
            if 'error' in item:
                self.error = True

        self.memory_ops.add_mapping(READ_POISON, (-1,))
        self.storage_ops.add_mapping(READ_POISON, (-1,))

    # self.__memory_ops.debug_contexts()

    def get_pair(self, opcode, index):
        if "0xfe" in opcode:
            opcode = "INVALID"
        delta, alpha = actions[opcode][:2]
        output = self.get_stack_items(index + 1, alpha)
        inputs = self.get_stack_items(index, delta)
        if opcode in call_ops:
            inputs = tuple([opcode] + inputs[1:])  # we do not care gas value
        else:
            inputs = tuple([opcode] + inputs)

        if len(output) == 1:
            output = output[0]
        else:
            output = None
        return inputs, output

    def get_stack_items(self, index, count=1):
        if count == 0:
            return []
        try:
            values = self.trace[index]['stack'][-count:]
        except IndexError:
            raise TraceException("trace ends with error")
        return [int(v, 16) for v in reversed(values)]

    # TODO: how to fix this stupid forwarding
    def fast_forward_trace(self, begin_address):
        begin = -1
        # print(begin_address)
        for index, item in enumerate(self.trace):
            opcode = item['op'].encode('utf-8')
            try:
                inputs, _ = self.get_pair(opcode, index)
            except TraceException:
                break
            if opcode in mem_write_ops:
                self.memory_ops.pop_context(inputs)

            if item['pc'] == begin_address:
                begin = index
                break

        if begin == -1 or len(self.trace) == 0:
            if self.error:
                raise TraceError("trace did not reach function entry")
            raise ValidationError("failed to reach entry")
        self.trace = self.trace[begin:]

    def get_cur_state(self):
        # print (self.trace[0])
        return self.trace[0]

    def do_mread(self, inputs):
        return self.memory_ops.lookup_mapping(inputs)

    def do_mwrite(self, inputs):
        self.memory_ops.pop_context(inputs)

    def do_sload(self, inputs):
        return self.storage_ops.lookup_mapping(inputs)

    def do_sstore(self, inputs):
        self.storage_ops.pop_context(inputs)

    def debug_mem_contexts(self):
        self.memory_ops.debug_mapping()


class EffectReader(TraceReader):
    def __init__(self, line):
        TraceReader.__init__(self, line)

        self.free_ops = FreeOps()
        self.storage_ops = ContextOps()
        self.memory_ops = dict()  # this is pretty bad
        self.effect_ops = EffectOps()

    def parse_trace(self):
        self.__parse_trace()

    # self.storage_ops.debug_mapping()

    def __parse_trace(self):
        for index, item in enumerate(self.trace):
            opcode = item['op'].encode('utf-8')
            # print(index)
            try:
                inputs, output = self.get_pair(opcode, index)
            except TraceException:
                self.trace = self.trace[:index]
                self.error = True
                break

            if opcode in free_ops:
                self.free_ops.add_mapping(inputs, output)
            elif opcode in effect_ops:
                self.__parse_effect_op(index, opcode, inputs, output)
            # if it still interacts with memory
            elif opcode in mem_read_ops | mem_write_ops:
                self.__parse_memory_ops(index, opcode, inputs, output)
            elif opcode == "SLOAD":
                self.storage_ops.add_mapping(inputs, output)
            elif opcode == "SSTORE":
                self.storage_ops.push_context(inputs)
            if 'error' in item:
                self.error = True
        # self.effect_ops.debug_mapping()
        self.storage_ops.add_mapping(READ_POISON, (-1,))

    def get_memory_bytes(self, index, offset, size):
        memory = self.trace[index]['memory']
        memory = "".join(memory)
        begin, end = offset * 2, (offset + size) * 2
        return memory[begin:end].encode('utf-8')

    def __parse_effect_op(self, index, opcode, inputs, output):
        if opcode in log_ops | {"RETURN"}:
            offset, size = inputs[1], inputs[2]
            chunk = self.get_memory_bytes(index, offset, size)
            inputs = tuple([opcode, chunk] + list(inputs[3:]))
            self.effect_ops.add_mapping(inputs, output)
        elif opcode in call_ops:  # gas is already removed
            # to [1], value [2], in offset [-4], in_size [-3], out_offset [-2], out_size [-1]
            in_offset, in_size = inputs[-4], inputs[-3]
            in_chunk = self.get_memory_bytes(index, in_offset, in_size)
            out_offset, out_size = inputs[-2], inputs[-1]
            out_chunk = self.get_memory_bytes(index + 1, out_offset, out_size)
            output = (output, out_chunk)
            if opcode == "DELEGATECALL":
                inputs = (opcode, inputs[1], in_chunk, out_offset, out_size)
            else:
                inputs = (opcode, inputs[1], inputs[2], in_chunk, out_offset, out_size)
            self.effect_ops.add_mapping(inputs, output)
        elif opcode == "CREATE":
            offset, size = inputs[2], inputs[3]
            chunk = self.get_memory_bytes(index, offset, size)
            inputs = (opcode, chunk)
            self.effect_ops.add_mapping(inputs, output)
        else:
            raise NotImplementedError("effect op " + opcode + "not parsed")

    def __parse_memory_ops(self, index, opcode, inputs, output):
        if opcode == "SHA3":
            offset, size = inputs[1], inputs[2]  # get the real input
            chunk = self.get_memory_bytes(index, offset, size)
            inputs = (opcode, chunk)
            # print(inputs)
            self.memory_ops[inputs] = output
        elif opcode == "CALLDATACOPY":
            offset, size = inputs[1], inputs[3]
            chunk = self.get_memory_bytes(index + 1, offset, size)
            self.memory_ops[inputs] = chunk
        elif opcode == "CODECOPY":
            offset, size = inputs[1], inputs[3]
            chunk = self.get_memory_bytes(index + 1, offset, size)
            self.memory_ops[inputs] = chunk
        elif opcode == "EXTCODECOPY":
            offset, size = inputs[2], inputs[4]
            chunk = self.get_memory_bytes(index + 1, offset, size)
            self.memory_ops[inputs] = chunk
        # print(self.trace[index+1])
        elif opcode not in {"MLOAD", "MSTORE", "MSTORE8"}:
            raise NotImplementedError("memory op " + opcode + " not parsed")

    def fast_forward_trace(self, begin_address):
        begin = -1
        for index, item in enumerate(self.trace):
            opcode = item['op'].encode('utf-8')
            try:
                inputs, _ = self.get_pair(opcode, index)
            except TraceException:
                break

            if item['pc'] == begin_address:
                begin = index
                break

        if begin == -1 or len(self.trace) == 0:
            if self.error:
                raise TraceError("trace did not reach function entry")
            raise ValidationError("failed to reach entry")
        self.trace = self.trace[begin:]

    def do_mem_ops(self, inputs):
        return self.memory_ops[inputs]

    def do_effect_ops(self, inputs):
        return self.effect_ops.lookup_mapping(inputs)

    def do_end_check(self):
        return self.effect_ops.reached_end_state()


if __name__ == "__main__":
    line = open(sys.argv[1]).readline()
    t = EffectReader(line)
    t.parse_trace()

# file = open(sys.argv[1])
# for count, line in enumerate(file):
# 	if count == 100:
# 		break
# 	try:
# 		print(count)
# 		t = EffectReader(line)
# 		t.parse_trace()
# 	except:
# 		pass