rossym plugin for ida

rossym-py3.py

'''
PROJECT:     ReactOS RosSym IDA Extension - python 3.x
LICENSE:     MIT (https://spdx.org/licenses/MIT)
PURPOSE:     Decode RosSym symbols
COPYRIGHT:   Copyright 2017-2019 Mark Jansen (mark.jansen@reactos.org)
'''
import idaapi
import idautils
import struct
import ctypes
import ida_nalt
import ida_name
import idc

class ROSSYM_HEADER(ctypes.LittleEndianStructure):
    _fields_ = [
        ("SymbolsOffset", ctypes.c_uint32),
        ("SymbolsLength", ctypes.c_uint32),
        ("StringsOffset", ctypes.c_uint32),
        ("StringsLength", ctypes.c_uint32),
    ]

class ROSSYM_ENTRY(ctypes.LittleEndianStructure):
    _fields_ = [
        ("Address",         ctypes.c_uint32),
        ("FunctionOffset",  ctypes.c_uint32),
        ("FileOffset",      ctypes.c_uint32),
        ("SourceLine",      ctypes.c_uint32),
    ]

def is_metadata(name):
    if len(name) < 3:
        return False
    return name[0] == '_' and name[1] != '_' and name[-1] == '_' and name[-2] == '_'

def read_struct(data, offset, struct):
    s = struct()
    slen = ctypes.sizeof(s)
    bytes = data[offset:offset+slen]
    ctypes.memmove(ctypes.addressof(s), bytes, slen)
    return s

def read_rossym(filename):
    try:
        with open(filename, 'rb') as f:
            if f.read(2) != b'MZ':
                print(filename, 'No dos header found!')
                return None
            f.seek(0x3C)
            e_lfanew = struct.unpack('i', f.read(4))[0]
            f.seek(e_lfanew)
            if f.read(4) != b'PE\0\0':
                print(filename, 'No PE header found!')
                return None
            f.seek(e_lfanew + 0x18)
            Magic = struct.unpack('h', f.read(2))[0]
            if Magic != 0x10b:
                print(filename, 'is not a 32 bit exe!')
                return None
            f.seek(e_lfanew + 0x6)
            NumberOfSections = struct.unpack('H', f.read(2))[0]
            f.seek(e_lfanew + 0x14)
            SizeOfOptionalHeader = struct.unpack('H', f.read(2))[0]
            FirstSection = e_lfanew + 0x18 + SizeOfOptionalHeader
            f.seek(FirstSection)
            for n in range(0, NumberOfSections):
                sect = f.read(0x28) # sizeof(IMAGE_SECTION_HEADER)
                Name, SizeOfRawData, PointerToRawData = struct.unpack('8s4x4x2i16x', sect)
                if Name == b'.rossym\0':
                    f.seek(PointerToRawData)
                    return f.read(SizeOfRawData)
    except IOError as ex:
        print('Error opening "{0}": ({1}, {2})'.format(filename, ex.errno, ex.strerror)) 
    return None

def ea2name(ea):
    return idaapi.get_name(ea)

def rename(ea, name):
    f = idaapi.get_flags(ea)
    if idaapi.has_user_name(f):
        #print 'Skipping', ea2name(ea), '(user name)'
        pass
    elif not idaapi.set_name(ea, name, idaapi.SN_NOCHECK | idaapi.SN_PUBLIC | idaapi.SN_NOWARN | idaapi.SN_FORCE):
        print('Failed renaming {:x}("{}") to {}'.format(ea, ea2name(ea), name))


class rossym_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Load ReactOS symbol info"

    help = "Load ReactOS symbol info"
    wanted_name = "RosSym"
    wanted_hotkey = "Ctrl-F8"

    def init(self):
        info = idaapi.get_inf_structure()
        if info.filetype != idaapi.f_PE:
            return idaapi.PLUGIN_SKIP

        self.rossym_data = read_rossym(ida_nalt.get_input_file_path())
        if self.rossym_data is None:
            return idaapi.PLUGIN_SKIP

        print(self.wanted_name, 'loaded.')
        return idaapi.PLUGIN_OK

    def run(self, arg):
        self.rossym = read_struct(self.rossym_data, 0, ROSSYM_HEADER)
        entry = ROSSYM_ENTRY()
        entry_size = ctypes.sizeof(entry)
        base = idaapi.get_imagebase()
        functions = {}
        renamed_ea = {}
        for sym in range(self.rossym.SymbolsOffset, self.rossym.SymbolsOffset + self.rossym.SymbolsLength, entry_size):
            ctypes.memmove(ctypes.addressof(entry), self.rossym_data[sym:sym+entry_size], entry_size)

            if entry.FunctionOffset not in functions:
                funcname = self.read_string(entry.FunctionOffset)
                functions[entry.FunctionOffset] = funcname
            else:
                funcname = functions[entry.FunctionOffset]
            if not is_metadata(funcname):
                address = base + entry.Address if funcname != '__ImageBase' and funcname != '__RUNTIME_PSEUDO_RELOC_LIST__' else entry.Address
                und = ida_name.demangle_name(funcname, 0)
                func = idaapi.get_func(address)
                if func:
                    if und:
                        #print(und)
                        idc.SetType(func.start_ea, und)
                    if func.start_ea not in renamed_ea:
                        renamed_ea[func.start_ea] = funcname
                        rename(func.start_ea, funcname)
                        #print funcname
                elif entry.FileOffset == 0:
                    if und:
                        idc.SetType(address, und)
                    rename(address, funcname)
                    

    def read_string(self, offset):
        str = ''
        while True:
            c = self.rossym_data[self.rossym.StringsOffset + offset]
            if c == 0:
                return str
            offset += 1
            str += chr(c)


    def term(self):
        pass


def PLUGIN_ENTRY():
    return rossym_t()

rossym_py2.py

'''
PROJECT:     ReactOS RosSym IDA Extension - python 2.x
LICENSE:     MIT (https://spdx.org/licenses/MIT)
PURPOSE:     Decode RosSym symbols
COPYRIGHT:   Copyright 2017,2018 Mark Jansen (mark.jansen@reactos.org)
'''
import idaapi
import idautils
import struct
import ctypes
import ida_nalt
import ida_name
import idc

class ROSSYM_HEADER(ctypes.LittleEndianStructure):
    _fields_ = [
        ("SymbolsOffset", ctypes.c_uint32),
        ("SymbolsLength", ctypes.c_uint32),
        ("StringsOffset", ctypes.c_uint32),
        ("StringsLength", ctypes.c_uint32),
    ]

class ROSSYM_ENTRY(ctypes.LittleEndianStructure):
    _fields_ = [
        ("Address",         ctypes.c_uint32),
        ("FunctionOffset",  ctypes.c_uint32),
        ("FileOffset",      ctypes.c_uint32),
        ("SourceLine",      ctypes.c_uint32),
    ]

def is_metadata(name):
    if len(name) < 3:
        return False
    return name[0] == '_' and name[1] != '_' and name[-1] == '_' and name[-2] == '_'

def read_struct(data, offset, struct):
    s = struct()
    slen = ctypes.sizeof(s)
    bytes = data[offset:offset+slen]
    ctypes.memmove(ctypes.addressof(s), bytes, slen)
    return s

def read_rossym(filename):
    try:
        with open(filename, 'rb') as f:
            if f.read(2) != 'MZ':
                print filename, 'No dos header found!'
                return None
            f.seek(0x3C)
            e_lfanew = struct.unpack('i', f.read(4))[0]
            f.seek(e_lfanew)
            if f.read(4) != 'PE\0\0':
                print filename, 'No PE header found!'
                return None
            f.seek(e_lfanew + 0x18)
            Magic = struct.unpack('h', f.read(2))[0]
            if Magic != 0x10b:
                print filename, 'is not a 32 bit exe!'
                return None
            f.seek(e_lfanew + 0x6)
            NumberOfSections = struct.unpack('H', f.read(2))[0]
            f.seek(e_lfanew + 0x14)
            SizeOfOptionalHeader = struct.unpack('H', f.read(2))[0]
            FirstSection = e_lfanew + 0x18 + SizeOfOptionalHeader
            f.seek(FirstSection)
            for n in xrange(0, NumberOfSections):
                sect = f.read(0x28) # sizeof(IMAGE_SECTION_HEADER)
                Name, SizeOfRawData, PointerToRawData = struct.unpack('8s4x4x2i16x', sect)
                if Name == '.rossym\0':
                    f.seek(PointerToRawData)
                    return f.read(SizeOfRawData)
    except IOError as ex:
        print 'Error opening "{0}": ({1}, {2})'.format(filename, ex.errno, ex.strerror) 
    return None

def ea2name(ea):
    return idaapi.get_name(ea)

def rename(ea, name):
    f = idaapi.get_flags(ea)
    if idaapi.has_user_name(f):
        #print 'Skipping', ea2name(ea), '(user name)'
        pass
    elif not idaapi.set_name(ea, name, idaapi.SN_NOCHECK | idaapi.SN_PUBLIC | idaapi.SN_NOWARN):
        print 'Failed renaming', ea, ':', ea2name(ea), 'to', name


class rossym_t(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "Load ReactOS symbol info"

    help = "Load ReactOS symbol info"
    wanted_name = "RosSym"
    wanted_hotkey = "Ctrl-F8"

    def init(self):
        info = idaapi.get_inf_structure()
        if info.filetype != idaapi.f_PE:
            return idaapi.PLUGIN_SKIP

        self.rossym_data = read_rossym(ida_nalt.get_input_file_path())
        if self.rossym_data is None:
            return idaapi.PLUGIN_SKIP

        print self.wanted_name, 'loaded.'
        return idaapi.PLUGIN_OK

    def run(self, arg):
        self.rossym = read_struct(self.rossym_data, 0, ROSSYM_HEADER)
        entry = ROSSYM_ENTRY()
        entry_size = ctypes.sizeof(entry)
        base = idaapi.get_imagebase()
        functions = {}
        renamed_ea = {}
        for sym in xrange(self.rossym.SymbolsOffset, self.rossym.SymbolsOffset + self.rossym.SymbolsLength, entry_size):
            ctypes.memmove(ctypes.addressof(entry), self.rossym_data[sym:sym+entry_size], entry_size)

            if entry.FunctionOffset not in functions:
                funcname = self.read_string(entry.FunctionOffset)
                functions[entry.FunctionOffset] = funcname
            else:
                funcname = functions[entry.FunctionOffset]
            if not is_metadata(funcname):
                address = base + entry.Address if funcname != '__ImageBase' and funcname != '__RUNTIME_PSEUDO_RELOC_LIST__' else entry.Address
                und = ida_name.demangle_name(funcname, 0)
                func = idaapi.get_func(address)
                if func:
                    if und:
                        print und
                        idc.SetType(func.start_ea, und)
                    if func.start_ea not in renamed_ea:
                        renamed_ea[func.start_ea] = funcname
                        rename(func.start_ea, funcname)
                        #print funcname
                elif entry.FileOffset == 0:
                    if und:
                        idc.SetType(address, und)
                    rename(address, funcname)
                    

    def read_string(self, offset):
        str = ''
        while True:
            c = self.rossym_data[self.rossym.StringsOffset + offset]
            if ord(c) == 0:
                return str
            offset += 1
            str += c


    def term(self):
        pass


def PLUGIN_ENTRY():
    return rossym_t()