asm_splitter.py

#what we need:

#.include "macros.inc"

#.section .text  # 0x80005940 - 0x803B7240

import sys, os, errno, re
import xlrd

# -============================================- #
#                  VARIABLES
# -============================================- #
args: list = sys.argv
textloc: str
textfile = None
maploc: str
undefSyms = []
symDict = {"":""}
linkerfile = None

xlsx: xlrd.book.Book
map: xlrd.sheet.Sheet

idx_address: int = 2
idx_name: int = 3
idx_ofile: int = 4
idx_libfile: int = 5
# -============================================- #
#                  FUNCTIONS
# -============================================- #

def GetLibFile(name: str) -> str:
    name = name.strip()

    ext = name.rfind('.')
    if (ext != -1):
        name = name[:ext]
    else:
        name = ""
    
    name = name.replace(".", "_")
    return name

def GetObjectFile(name: str) -> str:
    name = name.strip()
    endPath = name.rfind("\\")

    # strip file paths
    if (endPath != -1):
        name = name[endPath + 1:]

    # remove extension
    ext = name.rfind('.')
    if (ext != -1):
        name = name[:ext]
    orig = name

    # expand paths
    name = (name.replace("_", "/"))
    if (name.find("/") != -1):
        name = name[:name.rfind("/") + 1] + orig
    else:
        name = orig
    #print(name)
    return name

def GetAddress(address: str) -> str:
    address = address.strip()
    try:
        _ = int(address, 16)
    except:
        print("Invalid address found on line %s!" % line)
        exit()
    else:
        return address

def GetName(name: str) -> str:
    name = name.strip()

    if (name == "*fill*"):
        return ""

    template = name.find("<")
    if (template != -1):
        name = name[:template]

    proto = name.find("(")
    if (proto != -1):
        name = name[:proto]

    dtor = name.find("~")
    if (dtor != -1):
        name = name[1:]
        name = name + "_NS_dtor"

    name = name.replace("::", "_NS_")
    name = name.replace("$", "_")
    name = name.replace("@", "_")
    return name

def GatherSyms():
    for line in textfile:
        sym: str

        decglobal: int = line.find(".global ")

        if decglobal != -1:
            # slice out ".global "
            sym = line[decglobal + len(".global "):]

            if sym in undefSyms:
                continue

            undefSyms.append(sym)

            continue
        
        possibleSyms = []
        
        fn = line.find("func_")

        if fn != -1:
            possibleSyms.append(fn)

        for idx in possibleSyms:
            # slice out everything up to and after the symbol
            sym = line[idx:]

            if (sym.find(" ") != -1):
                sym = sym[:sym.find(" ")]

            if (sym.find(":") != -1):
                sym = sym[:sym.find(":")]

            if (sym.find("@") != -1):
                sym = sym[:sym.find("@")]

            if (sym.find("-") != -1):
                sym = sym[:sym.find("-")]

            if (sym.find("\n") != -1):
                sym = sym[:sym.find("\n")]
            
            if (sym in undefSyms):
                continue
                    
            undefSyms.append(sym)
            print("Found %s" % sym)
            break
    textfile.seek(0)


def CreateSymDict():
    print("Creating symbol dictionary...")
    for i in range(1, map.nrows - 1):
        _addr = map.cell_value(i, 2).upper()
        _name = GetName(map.cell_value(i, 3))

        isLabel: bool = True

        for sym in undefSyms:
            symaddr = sym[5:] # func_
            if (symaddr == _addr):
                overloadCount: int = 0
                while (_name in symDict.values()):
                    overloadCount += 1
                    if (not (_name + "_X" + str(overloadCount) + "_") in symDict.values()):
                        _name = _name + "_X" + str(overloadCount) + "_"
                        break
                symDict[sym] = _name
                print("%s -> %s" % (sym ,_name))
                isLabel = False
                break
        
        if isLabel:
            _temp = "lbl_" + _addr
            overloadCount: int = 0
            while (_name in symDict.values()):
                overloadCount += 1
                if (not (_name + "_X" + str(overloadCount) + "_") in symDict.values()):
                    _name = _name + "_X" + str(overloadCount) + "_"
                    break
            symDict[_temp] = _name
            print("%s -> %s" % (_temp ,_name))

    textfile.seek(0)

def ReplaceSymsNoSplit(name, newname):
    file = open(name, "r")
    newfile = open(newname, "w+")
    for asmline in file:

        possibleSyms = []
        fn = [i for i in range(len(asmline)) if asmline.startswith("func_", i)]
        lbl = [i for i in range(len(asmline)) if asmline.startswith("lbl_", i)]
        decglobal: int = asmline.find(".global ")
        possibleSyms += fn
        possibleSyms += lbl

        sym: str

        if decglobal != -1:
            # slice out ".global "
            sym = asmline[decglobal + len(".global "):]
            if sym in symDict:
                asmline.replace(sym, symDict[sym])

        for idx in possibleSyms:
            # slice out everything up to and after the symbol
            sym = asmline[idx:]
            if (sym.find(" ") != -1):
                sym = sym[:sym.find(" ")]

            if (sym.find(":") != -1):
                sym = sym[:sym.find(":")]

            if (sym.find("@") != -1):
                sym = sym[:sym.find("@")]

            if (sym.find("-") != -1):
                sym = sym[:sym.find("-")]

            if (sym.find("\n") != -1):
                sym = sym[:sym.find("\n")]
            if sym in symDict:
                asmline = asmline.replace(sym, symDict[sym])
            else:
                pass
                #print ("ERROR: Symbol [%s] is not defined!" % sym)
                #exit()

        newfile.write(asmline)
    newfile.close()
    file.close()

# -============================================- #
#                    MAIN
# -============================================- #

if (len(args) != 4):
    print("Usage: [asm file] [syms xlsx] [init file]")
    exit()
textloc = sys.argv[1]
maploc = sys.argv[2]
initloc = sys.argv[3]
xlsx = xlrd.open_workbook(maploc)
map = xlsx.sheet_by_index(0)
textfile = open(textloc, "r")
linkerfile = open("SPLIT_LINKER", "w+")
GatherSyms()
CreateSymDict()

# assume there is a header
line: int = 1

while (1):
    libfile: str = GetLibFile(map.cell_value(line, idx_libfile))
    ofile: str = GetObjectFile(map.cell_value(line, idx_ofile))
    path: str
    dir: str
    curfile = None
    addr: str = GetAddress(map.cell_value(line, idx_address))
    nextaddr: str
    name: str = GetName(map.cell_value(line, idx_name))
    
    if (line < map.nrows - 1):
        nextaddr = GetAddress(map.cell_value(line + 1, idx_address)) # this may misbehave if *fill* entries arent listed as having the same address as the proceeding function
    else:
        nextaddr = addr

    # continue since these have already been included as part of the last function
    if (name == ""):
        line += 1
        continue

    if libfile != "":
        path = libfile + "/" + ofile
    else:
        path = ofile
    path = os.path.dirname(textloc) + "/" + path

    # If the file does not exist then create the directory and open a file for reading
    if (os.path.exists(os.path.dirname(path)) == False):
        try:
            os.makedirs(os.path.dirname(path))
        except OSError as exc:
            if exc.errno == errno.EEXIST and os.path.isdir(path):
                pass
            else:
                raise

    if (os.path.exists(path + ".s")):
        curfile = open(path + ".s", "a")
    else:
        curfile = open(path + ".s", "w+")
        curfile.write(".include \"macros.inc\"\n\n")
        curfile.write(".section .text, \"ax\" # %s\n\n\n" % (addr))
        linkerfile.write( path + ".o;\n")
    
    while(1):
        prev: int = textfile.tell()
        asmline: str = textfile.readline()

        if (asmline.find(nextaddr.upper()) != -1):
            textfile.seek(prev)
            break
        else:
            possibleSyms = []
            fn = [i for i in range(len(asmline)) if asmline.startswith("func_", i)]
            lbl = [i for i in range(len(asmline)) if asmline.startswith("lbl_", i)]
            decglobal: int = asmline.find(".global ")
            possibleSyms += fn
            possibleSyms += lbl

            sym: str

            if decglobal != -1:
                # slice out ".global "
                sym = asmline[decglobal + len(".global "):]
                if sym in symDict:
                    asmline.replace(sym, symDict[sym])

            for idx in possibleSyms:
                # slice out everything up to and after the symbol
                sym = asmline[idx:]
                if (sym.find(" ") != -1):
                    sym = sym[:sym.find(" ")]

                if (sym.find(":") != -1):
                    sym = sym[:sym.find(":")]

                if (sym.find("@") != -1):
                    sym = sym[:sym.find("@")]

                if (sym.find("-") != -1):
                    sym = sym[:sym.find("-")]

                if (sym.find("\n") != -1):
                    sym = sym[:sym.find("\n")]
                if sym in symDict:
                    # since the original assembly did not pick up that this was a function we also need to add a .global for it.
                    if (sym.startswith("lbl_")):
                        curfile.write(".global %s\n" % symDict[sym])

                    asmline = asmline.replace(sym, symDict[sym])
                else:
                    pass
                    #print ("ERROR: Symbol [%s] is not defined!" % sym)
                    #exit()

            curfile.write(asmline)

    curfile.close()
    line += 1

    print(path + ": " + name)

    if line > map.nrows - 1:
        ReplaceSymsNoSplit(initloc, "init.s.new")
        break

textfile.close()
linkerfile.close()