def open_json_trace(filename):
    """Opens JSON trace file and reads trace data and bookmarks

    Args:
        filename: name of trace file
    """
    try:
        f = open(filename)
    except IOError:
        print("Error, could not open file.")
    else:
        with f:
            try:
                trace_data = TraceData()
                data = json.load(f)
                trace_data.bookmarks = []
                trace_data.filename = filename
                trace_data.trace = data["trace"]
                trace_data.arch = data.get("arch", "")
                trace_data.ip_reg = data.get("ip_reg", "")
                trace_data.pointer_size = data.get("pointer_size", 4)
                trace_data.regs = data.get("regs", {})
                if "bookmarks" in data:
                    for bookmark in data["bookmarks"]:
                        trace_data.add_bookmark(Bookmark(**bookmark))
                return trace_data
            except KeyError:
                print("Error while reading trace file.")
            except Exception as exc:
                print("Error while reading trace file: " + str(exc))
                print(traceback.format_exc())
    return None
def open_x64dbg_trace(filename):
    """Opens x64dbg trace file

    Args:
        filename: name of trace file
    Returns:
        TraceData object
    """
    with open(filename, "rb") as f:
        trace_data = TraceData()
        trace_data.filename = filename

        # check first 4 bytes
        magic = f.read(4)
        if magic != b"TRAC":
            raise ValueError("Error, wrong file format.")

        json_length_bytes = f.read(4)
        json_length = int.from_bytes(json_length_bytes, "little")

        # read JSON blob
        json_blob = f.read(json_length)
        json_str = str(json_blob, "utf-8")
        arch = json.loads(json_str)["arch"]

        reg_indexes = {}
        if arch == "x64":
            regs = prefs.X64_REGS
            ip_reg = "rip"
            capstone_mode = CS_MODE_64
            pointer_size = 8  # qword
        else:
            regs = prefs.X32_REGS
            ip_reg = "eip"
            capstone_mode = CS_MODE_32
            pointer_size = 4  # dword

        for i, reg in enumerate(regs):
            reg_indexes[reg] = i

        trace_data.arch = arch
        trace_data.ip_reg = ip_reg
        trace_data.regs = reg_indexes
        trace_data.pointer_size = pointer_size

        md = Cs(CS_ARCH_X86, capstone_mode)
        reg_values = [None] * len(reg_indexes)
        trace = []
        row_id = 0
        while f.read(1) == b"\x00":
            register_changes = int.from_bytes(f.read(1), "little")
            memory_accesses = int.from_bytes(f.read(1), "little")
            flags_and_opcode_size = int.from_bytes(f.read(1),
                                                   "little")  # Bitfield
            thread_id_bit = (flags_and_opcode_size >> 7) & 1  # msb
            opcode_size = flags_and_opcode_size & 15  # 4 lsbs

            if thread_id_bit > 0:
                thread_id = int.from_bytes(f.read(4), "little")

            opcodes = f.read(opcode_size)

            register_change_position = []
            for _ in range(register_changes):
                register_change_position.append(
                    int.from_bytes(f.read(1), "little"))

            register_change_new_data = []
            for _ in range(register_changes):
                register_change_new_data.append(
                    int.from_bytes(f.read(pointer_size), "little"))

            memory_access_flags = []
            for _ in range(memory_accesses):
                memory_access_flags.append(int.from_bytes(f.read(1), "little"))

            memory_access_addresses = []
            for _ in range(memory_accesses):
                memory_access_addresses.append(
                    int.from_bytes(f.read(pointer_size), "little"))

            memory_access_old_data = []
            for _ in range(memory_accesses):
                memory_access_old_data.append(
                    int.from_bytes(f.read(pointer_size), "little"))

            memory_access_new_data = []
            for i in range(memory_accesses):
                if memory_access_flags[i] & 1 == 0:
                    memory_access_new_data.append(
                        int.from_bytes(f.read(pointer_size), "little"))

            reg_id = 0
            for i, change in enumerate(register_change_position):
                reg_id += change
                if reg_id + i < len(reg_indexes):
                    reg_values[reg_id + i] = register_change_new_data[i]

            # disassemble
            ip_value = reg_values[reg_indexes[ip_reg]]
            for (_address, _size, mnemonic,
                 op_str) in md.disasm_lite(opcodes, ip_value):
                disasm = mnemonic
                if op_str:
                    disasm += " " + op_str

            mems = []
            mem = {}
            new_data_counter = 0
            for i in range(memory_accesses):
                flag = memory_access_flags[i]
                value = memory_access_old_data[i]
                mem["access"] = "READ"
                if flag & 1 == 0:
                    value = memory_access_new_data[new_data_counter]
                    mem["access"] = "WRITE"
                    new_data_counter += 1
                else:
                    pass
                    # memory value didn't change
                    # (it is read or overwritten with identical value)
                    # this has to be fixed somehow in x64dbg

                mem["addr"] = memory_access_addresses[i]

                # fix value (x64dbg saves all values as qwords)
                if "qword" in disasm:
                    pass
                elif "dword" in disasm:
                    value &= 0xFFFFFFFF
                elif "word" in disasm:
                    value &= 0xFFFF
                elif "byte" in disasm:
                    value &= 0xFF
                mem["value"] = value
                mems.append(mem.copy())

            trace_row = {}
            trace_row["id"] = row_id
            trace_row["ip"] = ip_value
            trace_row["disasm"] = disasm
            trace_row["regs"] = reg_values.copy()
            trace_row["opcodes"] = opcodes.hex()
            trace_row["mem"] = mems.copy()
            trace_row["comment"] = ""
            trace.append(trace_row)
            row_id += 1

        trace_data.trace = trace
        return trace_data
def open_tv_trace(filename):
    """Opens tvt trace file and reads trace data and bookmarks

    Args:
        filename: name of trace file
    """
    with open(filename, "rb") as f:
        trace_data = TraceData()
        trace_data.filename = filename

        # check first 4 bytes
        magic = f.read(4)
        if magic != b"TVTR":
            raise ValueError("Error, wrong file format.")

        json_length_bytes = f.read(4)
        json_length = int.from_bytes(json_length_bytes, "little")

        # read JSON blob
        json_blob = f.read(json_length)
        json_str = str(json_blob, "utf-8")
        file_info = json.loads(json_str)

        arch = file_info.get("arch", "")

        reg_indexes = {}
        if arch == "x64":
            regs = prefs.X64_REGS
            for i, reg in enumerate(regs):
                reg_indexes[reg] = i
            pointer_size = 8  # qword
            ip_reg = "rip"
        elif arch == "x86":
            regs = prefs.X32_REGS
            for i, reg in enumerate(regs):
                reg_indexes[reg] = i
            pointer_size = 4  # dword
            ip_reg = "eip"
        else:
            print(f"Unknown CPU arch: {arch} Let's try to load it anyway.")
            ip_reg = file_info.get("ip_reg", "")
            pointer_size = file_info.get("pointer_size", 4)

        if "regs" in file_info:
            reg_indexes = {}
            for i, reg in enumerate(file_info["regs"]):
                reg_indexes[reg] = i

        trace_data.arch = arch
        trace_data.ip_reg = ip_reg
        trace_data.regs = reg_indexes
        trace_data.pointer_size = pointer_size

        reg_values = [None] * len(reg_indexes)
        trace = []
        row_id = 0
        while f.peek(1)[:1] == b"\x00":
            f.read(1)
            disasm = ""
            disasm_length = int.from_bytes(f.read(1), "little")
            if disasm_length > 0:
                disasm = f.read(disasm_length).decode()

            comment = ""
            comment_length = int.from_bytes(f.read(1), "little")
            if comment_length > 0:
                comment = f.read(comment_length).decode()

            register_changes = int.from_bytes(f.read(1), "little")
            memory_accesses = int.from_bytes(f.read(1), "little")
            flags_and_opcode_size = int.from_bytes(f.read(1),
                                                   "little")  # Bitfield
            thread_id_bit = (flags_and_opcode_size >> 7) & 1  # msb
            opcode_size = flags_and_opcode_size & 15  # 4 lsb

            if thread_id_bit > 0:
                thread_id = int.from_bytes(f.read(4), "little")

            opcodes = f.read(opcode_size)

            register_change_position = []
            for _ in range(register_changes):
                register_change_position.append(
                    int.from_bytes(f.read(1), "little"))

            register_change_new_data = []
            for _ in range(register_changes):
                register_change_new_data.append(
                    int.from_bytes(f.read(pointer_size), "little"))

            memory_access_flags = []
            for _ in range(memory_accesses):
                memory_access_flags.append(int.from_bytes(f.read(1), "little"))

            memory_access_addresses = []
            for _ in range(memory_accesses):
                memory_access_addresses.append(
                    int.from_bytes(f.read(pointer_size), "little"))

            memory_access_data = []
            for i in range(memory_accesses):
                memory_access_data.append(
                    int.from_bytes(f.read(pointer_size), "little"))

            reg_id = 0
            for i, change_pos in enumerate(register_change_position):
                reg_id += change_pos
                if reg_id + i < len(reg_indexes):
                    reg_values[reg_id + i] = register_change_new_data[i]

            mems = []
            mem = {}
            for i in range(memory_accesses):
                flag = memory_access_flags[i]
                value = memory_access_data[i]
                mem["access"] = "READ"
                if flag & 1 == 1:
                    mem["access"] = "WRITE"

                mem["addr"] = memory_access_addresses[i]
                mem["value"] = value
                mems.append(mem.copy())

            trace_row = {}
            trace_row["id"] = row_id
            if ip_reg:
                trace_row["ip"] = reg_values[reg_indexes[ip_reg]]
            trace_row["disasm"] = disasm
            trace_row["comment"] = comment
            trace_row["regs"] = reg_values.copy()
            trace_row["opcodes"] = opcodes.hex()
            trace_row["mem"] = mems.copy()
            trace.append(trace_row)
            row_id += 1

        trace_data.trace = trace

        while f.peek(1)[:1] == b"\x01":
            f.read(1)
            bookmark = Bookmark()
            bookmark.startrow = int.from_bytes(f.read(4), "little")
            bookmark.endrow = int.from_bytes(f.read(4), "little")
            disasm_length = int.from_bytes(f.read(1), "little")
            bookmark.disasm = f.read(disasm_length).decode()
            comment_length = int.from_bytes(f.read(1), "little")
            bookmark.comment = f.read(comment_length).decode()
            addr_length = int.from_bytes(f.read(1), "little")
            bookmark.addr = f.read(addr_length).decode()
            trace_data.add_bookmark(bookmark)

        return trace_data