ps3.py

# Find and load exports, imports, and opd section for PS3 elf's.
#@author aerosoul
#@category Analysis
#@keybinding 
#@menupath 
#@toolbar 

# NOTE: 
#  You may also want to enable Propagate R2 Pointer for the PowerPC 
#  Constant Reference Analyzer in order to fill in the remaining 
#  functions with the r2 value.
#  Make sure that r2 is an unaffected register in your cspec file as
#  it was not in ppc_64_32.cspec.

# TODO: In the future, it's preferable to have a loader load OPD.
#  That way we can use the script solely for updating imports/exports.

from ghidra.program.model.data import ArrayDataType
from ghidra.program.model.data import CharDataType
from ghidra.program.model.data import StructureDataType
from ghidra.program.model.data import CategoryPath
from ghidra.program.model.data import DataTypeConflictHandler
from ghidra.program.model.data import UnsignedCharDataType
from ghidra.program.model.data import UnsignedShortDataType
from ghidra.program.model.data import UnsignedIntegerDataType
from ghidra.program.model.data import Pointer32DataType
from ghidra.program.model.data import StringDataType
from ghidra.program.model.symbol import SourceType
from ghidra.program.model.lang import RegisterValue
from ghidra.program.model.util import CodeUnitInsertionException
from ghidra.app.util.bin.StructConverter import *
from java.math import BigInteger
from java.lang import Exception as JavaException
import xml.etree.ElementTree as ET
import os
import sys


dataTypeManager = currentProgram.getDataTypeManager()
memory = currentProgram.getMemory()
listing = currentProgram.getListing()

# Load xml
scrdir = os.path.expanduser('~/ghidra_scripts')
tree = ET.parse(scrdir + '/ps3.xml')
root_xml = tree.getroot()


def forceDWord(addr):
    try:
        createDWord(addr)
    except CodeUnitInsertionException:
        clearListing(addr, addr.add(4))
        try:
            createDWord(addr)
        except:
            pass # idk, just don't end the script..
    except JavaException as e:
        print(e)

def forcePointer(addr):
    try:
        createData(addr, Pointer32DataType.dataType)
    except CodeUnitInsertionException:
        clearListing(addr, addr.add(4))
        try:
            createData(addr, Pointer32DataType.dataType)
        except:
            pass # idk, just don't end the script..
    except JavaException as e:
        print(e)
        

def createModInfoStructure():
    modInfoStruc = StructureDataType(CategoryPath.ROOT, "_scemoduleinfo", 0)
    modInfoStruc.add(WORD, 2, "modattribute", "")
    modInfoStruc.add(ArrayDataType(BYTE, 2, 1), 2, "modversion", "")
    modInfoStruc.add(ArrayDataType(CharDataType(), 27, 1), 27, "modname", "")
    modInfoStruc.add(CharDataType(), 1, "terminal", "")
    modInfoStruc.add(DWORD, 4, "gp_value", "")
    modInfoStruc.add(Pointer32DataType(), 4, "ent_top", "")
    modInfoStruc.add(Pointer32DataType(), 4, "ent_end", "")
    modInfoStruc.add(Pointer32DataType(), 4, "stub_top", "")
    modInfoStruc.add(Pointer32DataType(), 4, "stub_end", "")
    return modInfoStruc

def getImportStructure():
    libstubstruc = StructureDataType(CategoryPath.ROOT, "_scelibstub", 0)
    libstubstruc.add(BYTE, 1, "structsize", "")
    libstubstruc.add(BYTE, 1, "reserved1", "")
    libstubstruc.add(WORD, 2, "version", "")
    libstubstruc.add(WORD, 2, "attribute", "")
    libstubstruc.add(WORD, 2, "nfunc", "")
    libstubstruc.add(WORD, 2, "nvar", "")
    libstubstruc.add(WORD, 2, "ntlsvar", "")
    libstubstruc.add(ArrayDataType(BYTE, 4, 1), 4, "reserved2", "")
    libstubstruc.add(Pointer32DataType(), 4, "libname", "")
    libstubstruc.add(Pointer32DataType(), 4, "func_nidtable", "")
    libstubstruc.add(Pointer32DataType(), 4, "func_table", "")
    libstubstruc.add(Pointer32DataType(), 4, "var_nidtable", "")
    libstubstruc.add(Pointer32DataType(), 4, "var_table", "")
    libstubstruc.add(Pointer32DataType(), 4, "tls_nidtable", "")
    libstubstruc.add(Pointer32DataType(), 4, "tls_table", "")
    return libstubstruc

def getExportStructure():
    libentstruc = StructureDataType(CategoryPath.ROOT, "_scelibent", 0)
    libentstruc.add(BYTE, 1, "structsize", "")
    libentstruc.add(BYTE, 1, "reserved", "")
    libentstruc.add(WORD, 2, "version", "")
    libentstruc.add(WORD, 2, "attribute", "")
    libentstruc.add(WORD, 2, "nfunc", "")
    libentstruc.add(WORD, 2, "nvar", "")
    libentstruc.add(WORD, 2, "ntlsvar", "")
    libentstruc.add(ArrayDataType(BYTE, 4, 1), 4, "reserved", "")
    libentstruc.add(Pointer32DataType(), 4, "libname", "")
    libentstruc.add(Pointer32DataType(), 4, "nidtable", "")
    libentstruc.add(Pointer32DataType(), 4, "addtable", "")
    return libentstruc



def get_name_for_nid(nid, library):
    for group in root_xml:
        #print "%s - %s" % (group.attrib['name'], library)
        if group.attrib['name'] == library:
            for entry in group:
                if int(entry.attrib['id'], 16) == nid:
                    return entry.attrib['name']
                    
def loadExports(start, end):
    libentstruc = getExportStructure()
    dataTypeManager.addDataType(libentstruc, DataTypeConflictHandler.REPLACE_HANDLER)

    addr = start
    print "Loading exports %s.." % addr
    while addr < end:
        #print addr
        
        nfunc = memory.getShort(addr.add(0x6))
        
        func_libname_ptr = memory.getInt(addr.add(0x10))
        func_nidtable = toAddr( memory.getInt(addr.add(0x14)) )
        func_table    = toAddr( memory.getInt(addr.add(0x18)) )
            
        if func_libname_ptr == 0:
            libname = ""
        else:
            libname = listing.getDataAt(toAddr(func_libname_ptr)).getDefaultValueRepresentation().encode('utf-8').strip("\"")
            try:
                listing.createData(toAddr(func_libname_ptr), StringDataType())
            except:
                pass
        #print libname
        
        for x in xrange(nfunc):
            func_nid  = memory.getInt(func_nidtable) & 0xffffffff
            func_ptr  = toAddr(memory.getInt(func_table))
            

            forceDWord(func_nidtable)
            forcePointer(func_table)

            func_name = get_name_for_nid(func_nid, libname)
            #print func_name
            if func_name != None:
                print "%s %s" % (libname, func_name)
                setEOLComment(func_nidtable, func_name)
                setEOLComment(func_table, hex(func_nid))
                opd_func_ptr = toAddr( memory.getInt(func_ptr) )
                func = listing.getFunctionAt(opd_func_ptr)
                if func != None:
                    func.setName(func_name, SourceType.ANALYSIS)
                else:
                    createFunction(opd_func_ptr, func_name)
            else:
                setEOLComment(func_nidtable, "?")
                setEOLComment(func_table, hex(func_nid))

            func_nidtable = func_nidtable.add(0x4)
            func_table    = func_table.add(0x4)
            

        try:
            createData(addr, libentstruc)
        except:
            pass
        
        addr = addr.add(0x1C)

def loadImports(start, end):
    # register structure into data type library
    libstubstruc = getImportStructure()
    dataTypeManager.addDataType(libstubstruc, DataTypeConflictHandler.REPLACE_HANDLER)
    
    addr = start
    print "Loading imports at %s.." % addr
    while addr < end:
        #print addr

        nfunc = memory.getShort(addr.add(0x6))

        func_libname_ptr = memory.getInt(addr.add(0x10))
        func_nidtable = toAddr( memory.getInt(addr.add(0x14)) )
        func_table    = toAddr( memory.getInt(addr.add(0x18)) )
            
        if func_libname_ptr == 0:
            libname = ""
        else:
            libname = listing.getDataAt(toAddr(func_libname_ptr)).getDefaultValueRepresentation().encode('utf-8').strip("\"")
            try:
                listing.createData(toAddr(func_libname_ptr), StringDataType())
            except:
                pass

        for x in xrange(nfunc):
            func_nid  = memory.getInt(func_nidtable) & 0xffffffff
            func_ptr  = toAddr(memory.getInt(func_table))

            forceDWord(func_nidtable)
            forcePointer(func_table)

            func_name = get_name_for_nid(func_nid, libname)
            #print func_name
            if func_name != None:
                print "%s %s" % (libname, func_name)
                setEOLComment(func_nidtable, func_name)
                setEOLComment(func_table, hex(func_nid))
                func = listing.getFunctionAt(func_ptr)
                if func != None:
                    func.setName(func_name, SourceType.ANALYSIS)
                else:
                    createFunction(func_ptr, func_name)
            else:
                setEOLComment(func_nidtable, "?")
                setEOLComment(func_table, hex(func_nid))

            func_nidtable = func_nidtable.add(0x4)
            func_table    = func_table.add(0x4)
            

        try:
            listing.createData(addr, libstubstruc)
        except:
            pass
        
        addr = addr.add(0x2C)
   

def getFunction(addr):
    # check if function exists
    func = getFunctionAt(addr)
    if func == None:
        clearListing(addr)
        disassemble(addr)
        return createFunction(addr, None)
    else:
        return func
        
def createOPDEntry(addr):
    func_ptr = toAddr( memory.getInt(addr) )
    func_toc = memory.getInt(addr.add(0x4))
    func = getFunction(func_ptr)
    
    #createFunction(func_ptr, None)
        
    func_min = func.getBody().getMinAddress()
    func_max = func.getBody().getMaxAddress()
    
    forcePointer(addr)
    forceDWord(addr.add(4))
    
    r2reg = currentProgram.getRegister("r2")
    tocValue = RegisterValue(r2reg, BigInteger.valueOf(func_toc))
    currentProgram.getProgramContext().setRegisterValue(func_min, func_max, tocValue)
        
def loadOPD():
    # get header segment
    # Block 0 is the ELF header, at offset 0x18 is e_entry.
    entry_point = memory.getLong(memory.getBlocks()[0].getStart().add(0x18))
    #print "%x" % entry_point
    # Get the opd memory block
    opd_sect = memory.getBlock(toAddr(entry_point))
    addr = opd_sect.getStart()
    print "Loading OPD at %s.." % addr
    while addr < opd_sect.getEnd():
        createOPDEntry(addr)
        addr = addr.add(0x8)
        
def findImports():
    for memBlock in memory.getBlocks():
        if not memBlock.isInitialized():
            continue
        sectaddr = memBlock.getStart()
        structsize = memBlock.getByte(sectaddr) & 0xff
        if structsize > memBlock.getSize():
            continue
        structsize2 = memBlock.getByte(sectaddr.add(structsize)) & 0xff
        if structsize == structsize2:
            if structsize == 0x2C:
                loadImports(memBlock.getStart(), memBlock.getEnd())
                break
                
def findExports():
    for memBlock in memory.getBlocks():
        if not memBlock.isInitialized():
            continue
        sectaddr = memBlock.getStart()
        structsize = memBlock.getByte(sectaddr) & 0xff
        if structsize > memBlock.getSize():
            continue
        structsize2 = memBlock.getByte(sectaddr.add(structsize)) & 0xff
        if structsize == structsize2:
            if structsize == 0x1C:
                loadExports(memBlock.getStart(), memBlock.getEnd())
                break
                
def isPrxElfType():
    # Elf header is not included in memory space for prx's
    # so _elfHeader memory block is created for it
    if getMemoryBlock('_elfHeader') != None:
        return True
        

def loadModuleInfo():
    progHeaders = getMemoryBlock('_elfProgramHeaders').getStart()
    p_paddr = getLong(progHeaders.add(0x18))
    p_offset = getLong(progHeaders.add(0x8))
    modInfoAddr = toAddr(p_paddr - p_offset)
    print("ModInfoAddr: {0}".format(modInfoAddr))
    modInfoStruc = createModInfoStructure()
    dataTypeManager.addDataType(modInfoStruc, DataTypeConflictHandler.REPLACE_HANDLER)
    createData(modInfoAddr, modInfoStruc)
    ent_top = toAddr( getInt(modInfoAddr.add(0x24)) )
    ent_end = toAddr( getInt(modInfoAddr.add(0x28)) )
    stub_top = toAddr( getInt(modInfoAddr.add(0x2C)) ) 
    stub_end = toAddr( getInt(modInfoAddr.add(0x30)) )
    loadExports(ent_top, ent_end)
    loadImports(stub_top, stub_end)
    gp_value = getInt(modInfoAddr.add(0x20))
    opdBlock = getMemoryBlocks()[1]
    addr = opdBlock.getStart()
    while addr < opdBlock.getEnd():
        tmp = getInt(addr)
        if tmp == gp_value:
            start = addr.subtract(4)
            print("Found OPD: {0}".format(start))
            break
        addr = addr.add(4)
    addr = start
    print("Creating OPD")
    tmp = getInt(addr.add(4))
    print(hex(tmp))
    print(hex(gp_value))
    while tmp == gp_value:
        createOPDEntry(addr)
        tmp = getInt(addr.add(4))
        addr = addr.add(8)
    

if __name__ == '__main__':
    if isPrxElfType():
        loadModuleInfo()
    else:
        findImports()
        findExports()  
        loadOPD()