'''
PROJECT:     ReactOS RosSym IDA Extension - python 3.x
LICENSE:     MIT (https://spdx.org/licenses/MIT)
PURPOSE:     Decode RosSym symbols
COPYRIGHT:   Copyright 2020 Mark Jansen (mark.jansen@reactos.org)
'''
try:
    import idaapi
    import idautils
    import ida_nalt
    import ida_name
    import idc
    IS_IDA = True
except ImportError:
    IS_IDA = False

import struct
import ctypes

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

if IS_IDA:
    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()


def dump(filename):
    print('Dumping:', filename)
    rossym_data = read_rossym(filename)
    if not rossym_data:
        print('No rossym found')
        return

    rossym = read_struct(rossym_data, 0, ROSSYM_HEADER)

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


    rossym = read_struct(rossym_data, 0, ROSSYM_HEADER)
    entry = ROSSYM_ENTRY()
    entry_size = ctypes.sizeof(entry)
    functions = {}
    for sym in range(rossym.SymbolsOffset, rossym.SymbolsOffset + rossym.SymbolsLength, entry_size):
        ctypes.memmove(ctypes.addressof(entry), rossym_data[sym:sym+entry_size], entry_size)

        funcname = read_string(entry.FunctionOffset)

        if not funcname in functions:
            functions[funcname] = [entry.Address]
        else:
            functions[funcname].append(entry.Address)

    for funcname in sorted(functions.keys()):
        if not is_metadata(funcname):
            print(funcname, len(functions[funcname]))


if __name__ == '__main__':
    import sys
    for arg in sys.argv[1:]:
        dump(arg)