# How to Use:
# Change the string below to point to the .map file generated by the recompiler
# Open the JIT code dump in IDA
# Go to File -> Script File and pick this script.

file_to_open = "D:/IDA/DynarecOutput.map" # **PATH TO YOUR MAP FILE GOES HERE**

import sys
from idautils import *
from idaapi import *

codeSegmentNumber = 0

def parse_symbol_line(line): # Parse symbols and add them to the db
    words = line.split()
    address = int(words[0], 16) # Symbol address
    symbolName = words[1]       # Symbol name

    if (idc.set_name(address, symbolName) == 0): # Check if the set_name call failed, try to work around it (needed for calls to C++ functions)
        print("Failed to set", symbolName, "symbol")
        global codeSegmentNumber
        idaapi.add_segm(0, address, address + 0x4000, "code" + str(codeSegmentNumber), ".code") # Create segment
        idc.set_name(address, symbolName) # Attempt to name the address again
        codeSegmentNumber += 1
    
    if (len(words) >= 3):
        size = int(words[2]) # If available, the size of the symbol (for variables)
        if (size == 1):
            idaapi.apply_tinfo(address, tinfo_t(BTF_UINT8), idaapi.TINFO_DEFINITE)
        elif (size == 2):
            idaapi.apply_tinfo(address, tinfo_t(BTF_UINT16), idaapi.TINFO_DEFINITE)
        elif (size == 4):
            idaapi.apply_tinfo(address, tinfo_t(BTF_UINT32), idaapi.TINFO_DEFINITE)
        elif (size == 8):
            idaapi.apply_tinfo(address, tinfo_t(BTF_UINT64), idaapi.TINFO_DEFINITE)
        elif (size == 16):
            idaapi.apply_tinfo(address, tinfo_t(BTF_UINT128), idaapi.TINFO_DEFINITE)

def extract_segments(file):
    codeCacheSegment = int(file.readline().split()[0], 16) # Read the base of the code buffer (first line of file)
    rebase_program(codeCacheSegment, MSF_FIXONCE) # "Move" the code segment to the proper address

    for line in file: # Read segment info
        if line.strip() == "endsegs()":
            return
        else:
            words = line.split()
            address = int(words[0], 16) # Base segment address (hex)
            name = words[1]             # Segment name
            size = int(words[2], 16)    # Segment size (hex)
            type = words[3]             # .data, .text, etc

            idaapi.add_segm(0, address, address + size, name, type) # Create segment
            print (f"Created {name} segment at {hex(address)}-{hex(address + size)} ({type})")

with open(file_to_open, "r") as file:
    extract_segments(file)

    for line in file: # load symbols
        parse_symbol_line(line)

print ("Complete")