MS Symbol helper (wrapper around dbghelp)

mssymbol.py

from ctypes import *
from ctypes import wintypes
from collections import namedtuple
import os
import tempfile

# CONFIG SYMBOLS
SYMBOLS_PATH = os.environ.get('_NT_SYMBOL_PATH')
if not SYMBOLS_PATH:
    CACHE_PDB_PATH = os.path.join(tempfile.gettempdir(), "symbols")
    try:
        os.mkdir(CACHE_PDB_PATH)
    except:
        pass
    SYMBOLS_PATH = "SRV*" + CACHE_PDB_PATH + "*http://msdl.microsoft.com/download/symbols"

# print("Using SYMBOLS PATH: {}".format(SYMBOLS_PATH))

# CONST / TYPES
MAX_PATH = 260
RPC_MAX_LENGTH = 260
MAX_SYM_NAME = 150
UINT    = c_uint
UINT64  = c_uint64
BOOL    = c_bool
DWORD   = wintypes.DWORD
ULONG   = c_ulong
USHORT  = c_ushort
WCHAR   = c_wchar
HANDLE  = wintypes.HANDLE
HICON   = HANDLE
USHORT_PTR = POINTER(USHORT)
ULONG_PTR = POINTER(c_ulong)
FILETIME = wintypes.FILETIME
windll.kernel32.GetCurrentProcess.restype = HANDLE
PROCESS_VM_READ = 0x0010
PROCESS_QUERY_INFORMATION = 0x0400
SYMFLAG_VALUEPRESENT     = 0x00000001
SYMFLAG_REGISTER         = 0x00000008
SYMFLAG_REGREL           = 0x00000010
SYMFLAG_FRAMEREL         = 0x00000020
SYMFLAG_PARAMETER        = 0x00000040
SYMFLAG_LOCAL            = 0x00000080
SYMFLAG_CONSTANT         = 0x00000100
SYMFLAG_EXPORT           = 0x00000200
SYMFLAG_FORWARDER        = 0x00000400
SYMFLAG_FUNCTION         = 0x00000800
SYMFLAG_VIRTUAL          = 0x00001000
SYMFLAG_THUNK            = 0x00002000
SYMFLAG_TLSREL           = 0x00004000
SYMFLAG_SLOT             = 0x00008000
SYMFLAG_ILREL            = 0x00010000
SYMFLAG_METADATA         = 0x00020000
SYMFLAG_CLR_TOKEN        = 0x00040000

class SymOpt:
    CASE_INSENSITIVE = 0x00000001
    UNDNAME = 0x00000002
    DEFERRED_LOADS = 0x00000004
    DEBUG = 0x80000000

# IMAGEHLP_SYMBOL_TYPE_INFO
TI_GET_SYMTAG     = 0
TI_GET_SYMNAME    = 1
TI_GET_LENGTH     = 2
TI_GET_TYPE       = 3
TI_GET_BASETYPE   = 5 
TI_GET_ARRAYINDEXTYPEID = 6
TI_FINDCHILDREN   = 7
TI_GET_OFFSET     = 10
TI_GET_COUNT      = 12
TI_GET_CHILDRENCOUNT = 13
TI_GET_UDTKIND = 24

# SymTagEnum 
SymTagUDT = 11
SymTagFunctionType = 13
SymTagPointerType = 14
SymTagArrayType = 15
SymTagBaseType = 16

BaseType = {
    1: 'void',
    2: 'char',
    3: 'wchar',
    6: 'int',
    7: 'uint',
    8: 'float',
    10: 'bool', 
    13: 'long',
    14: 'ulong'
}

UdtType = {
    0: 'struct ',
    1: 'class ',
    2: 'union ',
}

class SYMBOL_INFO(Structure): # _SYMBOL_INFO
    _fields_ = [("SizeOfStruct", ULONG),
                ("TypeIndex", ULONG),
                ("Reserved", UINT64 * 2),
                ("Index", ULONG),
                ("Size", ULONG),
                ("ModBase", UINT64),
                ("Flags", ULONG),
                ("Value", UINT64),
                ("Address", UINT64),
                ("Register", ULONG),
                ("Scope", ULONG),
                ("Tag", ULONG),
                ("NameLen", ULONG),
                ("MaxNameLen", ULONG),
                ("Name", c_char * MAX_SYM_NAME)]

Symbol = namedtuple('Symbol', ['Name', 'TypeIndex', 'Size', 'RVA', 'Elements'])
Element = namedtuple('Element', ['Name', 'TypeIndex', 'Offset', 'Type'])

SYM_ENUMERATESYMBOLS_CALLBACK = WINFUNCTYPE(BOOL, POINTER(SYMBOL_INFO), ULONG, c_void_p)

def find_file(filename):
    for directory in [os.getcwd()] + os.environ['path'].split(';'):
        if os.path.isfile(os.path.join(directory, filename)):
            return os.path.join(directory, filename)

def find_element(structure, name):
    return [element for element in structure.Elements if element.Name == name][0]

class MSSymbolHelper(object):
    instance = None
    def __new__(cls): # __new__ always a classmethod
        if not MSSymbolHelper.instance:
            MSSymbolHelper.instance = _Singleton_MSSymbolHelper()
        return MSSymbolHelper.instance
    def __getattr__(self, name):
        return getattr(self.instance, name)
    def __setattr__(self, name):
        return setattr(self.instance, name)

class _Singleton_MSSymbolHelper:
    def __init__(self):
        self.current_process = windll.kernel32.GetCurrentProcess()
        self.sym_init()
        self.current_loaded_mod = None
        self.current_loaded_base = None
    
    def sym_init(self):
        if not os.path.isfile(r'C:\Windows\System32\symsrv.dll'):
            # TODO Update symsrv must be in same directory than dbghelp
            # https://docs.microsoft.com/en-us/windows/win32/debug/using-symsrv#installation
            raw_input("Can't find SYMSRV.DLL in System32. Microsoft symbols won't be available ! Continue ?")
        self.dbghelp = windll.LoadLibrary("dbghelp.dll")
        self.dbghelp.SymSetOptions.argtypes = [DWORD]
        self.dbghelp.SymInitialize.argtypes = [HANDLE, c_char_p, BOOL]
        self.dbghelp.SymLoadModule64.argtypes = [HANDLE, HANDLE, c_char_p, c_char_p, UINT64, DWORD]
        self.dbghelp.SymLoadModule64.restype = UINT64
        self.dbghelp.SymUnloadModule64.argtypes = [HANDLE, UINT64]
        self.dbghelp.SymFromAddr.argtypes = [HANDLE, UINT64, POINTER(UINT64), POINTER(SYMBOL_INFO)]
        self.dbghelp.SymFromName.argtypes = [HANDLE, c_char_p, POINTER(SYMBOL_INFO)]
        self.dbghelp.SymEnumSymbols.argtypes = [HANDLE, UINT64, c_char_p, SYM_ENUMERATESYMBOLS_CALLBACK, c_void_p]
        self.dbghelp.SymEnumTypes.argtypes = [HANDLE, UINT64, SYM_ENUMERATESYMBOLS_CALLBACK, c_void_p]
        self.dbghelp.SymGetTypeFromName.argtypes = [HANDLE, UINT64, c_char_p, POINTER(SYMBOL_INFO)]
        self.dbghelp.SymGetTypeInfo.argtypes = [HANDLE, UINT64, ULONG, ULONG, c_void_p]
        self.dbghelp.SymSetOptions(SymOpt.CASE_INSENSITIVE | SymOpt.UNDNAME)
        self.dbghelp.SymInitialize(self.current_process, SYMBOLS_PATH, False)
        
    def util_init_symbol_info(self):
        symbol = SYMBOL_INFO()
        symbol.SizeOfStruct = ((SYMBOL_INFO.Name.offset / sizeof(c_void_p)) + 1) * sizeof(c_void_p) # Align
        assert(symbol.SizeOfStruct == 88)
        symbol.MaxNameLen = ULONG(MAX_SYM_NAME)
        return pointer(symbol)
    
    def sym_load_module(self, modpath):
        if '\\' not in modpath:
            modpath = find_file(modpath)
        if not self.current_loaded_mod or self.current_loaded_mod.lower() != modpath.lower():
            if self.current_loaded_base:
                self.dbghelp.SymUnloadModule64(self.current_process, self.current_loaded_base)
            self.current_loaded_base = self.dbghelp.SymLoadModule64(self.current_process, 0, modpath, None, 0, 0)
            self.current_loaded_mod = modpath
        return self.current_loaded_base
        
    def sym_from_addr(self, modpath, rva):
        hbase = self.sym_load_module(modpath)
        p_symbol = self.util_init_symbol_info()
        displacement = UINT64(0)
        if self.dbghelp.SymFromAddr(self.current_process, hbase + rva, displacement, p_symbol):
            append = ' + 0x{:x}'.format(displacement.value) if displacement else ''
            return p_symbol.contents.Name + append
        return None
        
    def sym_from_name(self, modpath, name):
        hbase = self.sym_load_module(modpath)
        p_symbol = self.util_init_symbol_info()
        if self.dbghelp.SymFromName(self.current_process, name, p_symbol):
            return p_symbol.contents.Address - hbase
        return None
    
    def _enum_callback(self, p_symbol, size, context):
        self.symbols[p_symbol.contents.Name] = Symbol(
            p_symbol.contents.Name, 
            p_symbol.contents.TypeIndex, 
            p_symbol.contents.Size,
            p_symbol.contents.Address - p_symbol.contents.ModBase,
            []
        )
        return 1

    def enum(self, modpath):
        hbase = self.sym_load_module(modpath)
        self.symbols = {}
        self.dbghelp.SymEnumSymbols(self.current_process, hbase, None, SYM_ENUMERATESYMBOLS_CALLBACK(self._enum_callback), None)
        return self.symbols

    def resolve(self, modpath, rva_or_name=None):
        if isinstance(rva_or_name, basestring):
            return self.sym_from_name(modpath, rva_or_name)
        else:
            return self.sym_from_addr(modpath, rva_or_name)
    
    def enum_types(self, modpath):
        hbase = self.sym_load_module(modpath)
        self.symbols = {}
        self.dbghelp.SymEnumTypes(self.current_process, hbase, SYM_ENUMERATESYMBOLS_CALLBACK(self._enum_callback), None)
        return self.symbols

    def _get_type_info(self, hbase, typeid):
        typeName = c_wchar_p()
        tag = DWORD()
        self.dbghelp.SymGetTypeInfo(self.current_process, hbase, typeid, TI_GET_SYMTAG, pointer(tag))
        typeName.value = str(tag.value)
        if tag.value == SymTagBaseType:
            basetype = DWORD()
            self.dbghelp.SymGetTypeInfo(self.current_process, hbase, typeid, TI_GET_BASETYPE, pointer(basetype))
            typeName.value = BaseType[basetype.value]
            length = UINT64()
            self.dbghelp.SymGetTypeInfo(self.current_process, hbase, typeid, TI_GET_LENGTH, pointer(length))
            if length:
                typeName.value += str(8*length.value)
        elif tag.value == SymTagPointerType:
            pointertype = DWORD()
            self.dbghelp.SymGetTypeInfo(self.current_process, hbase, typeid, TI_GET_TYPE, pointer(pointertype))
            return self._get_type_info(hbase, pointertype) + '*'
        elif tag.value == SymTagUDT:
            name = c_wchar_p()
            udt = DWORD()
            self.dbghelp.SymGetTypeInfo(self.current_process, hbase, typeid, TI_GET_UDTKIND, pointer(udt))
            self.dbghelp.SymGetTypeInfo(self.current_process, hbase, typeid, TI_GET_SYMNAME, pointer(name))
            return UdtType[udt.value] + name.value.encode('ascii')
        elif tag.value == SymTagFunctionType:
            return 'fct' # TODO type of args
        elif tag.value == SymTagArrayType:
            count = DWORD()
            arraytype = DWORD()
            self.dbghelp.SymGetTypeInfo(self.current_process, hbase, typeid, TI_GET_COUNT, pointer(count))
            self.dbghelp.SymGetTypeInfo(self.current_process, hbase, typeid, TI_GET_TYPE, pointer(arraytype))
            return self._get_type_info(hbase, arraytype) + '[{}]'.format(count.value)
        else:
            raise ValueError("Unknown tag {}".format(tag.value))
        return typeName.value.encode('ascii')

    def struct(self, modpath, name):
        hbase = self.sym_load_module(modpath)
        p_symbol = self.util_init_symbol_info()
        if self.dbghelp.SymGetTypeFromName(self.current_process, hbase, name, p_symbol):
            type_index = p_symbol.contents.TypeIndex
            struct = Symbol(
                p_symbol.contents.Name, 
                p_symbol.contents.TypeIndex, 
                p_symbol.contents.Size,
                p_symbol.contents.Address - p_symbol.contents.ModBase, 
                []
            )
            count = DWORD()
            if self.dbghelp.SymGetTypeInfo(self.current_process, hbase, type_index, TI_GET_CHILDRENCOUNT, pointer(count)):
                elements = (ULONG * (count.value + 2))()
                elements[0] = count
                if self.dbghelp.SymGetTypeInfo(self.current_process, hbase, type_index, TI_FINDCHILDREN, pointer(elements)):
                    for child in elements[2:]:
                        name = c_wchar_p()
                        offset = DWORD()
                        typeId = DWORD()
                        self.dbghelp.SymGetTypeInfo(self.current_process, hbase, child, TI_GET_SYMNAME, pointer(name))
                        self.dbghelp.SymGetTypeInfo(self.current_process, hbase, child, TI_GET_OFFSET, pointer(offset))
                        self.dbghelp.SymGetTypeInfo(self.current_process, hbase, child, TI_GET_TYPE, pointer(typeId))
                        typeStr = self._get_type_info(hbase, typeId)
                        struct.Elements.append(Element(name.value.encode('ascii'), child, offset.value, typeStr))
            return struct
        return None

if __name__ == '__main__':
    symbols = MSSymbolHelper()
    all_syms = symbols.enum('ntdll.dll')
    print("NTDLL!* length: {}".format(len(all_syms)))
    print("NTDLL!A_SHAFinal: 0x{:X}".format(all_syms['A_SHAFinal'].RVA))
    print("NTDLL32!A_SHAFinal: 0x{:X}".format(symbols.resolve(r'C:\Windows\SYSWOW64\ntdll.dll', 'A_SHAFinal')))
    print("NTDLL+0x1234: {}".format(symbols.resolve(r'ntdll.dll', 0x1234)))

    print("ntoskrnl!_DRIVER_OBJECT DeviceObject.")
    el = find_element(symbols.struct("ntoskrnl.exe", '_DRIVER_OBJECT'), 'DeviceObject')
    print("\t+0x{:x}\t{:20s}\t\t: [{}]".format(el.Offset, el.Name, el.Type))

    print("ntoskrnl!_KPROCESS")
    k_process =symbols.struct("ntoskrnl.exe", '_KPROCESS')
    for el in k_process.Elements[:4]:
        print("\t+0x{:x}\t{:20s}\t\t: [{}]".format(el.Offset, el.Name, el.Type))
    print("\t...")

""" Result with public symbols
NTDLL!* length: 5334
NTDLL!A_SHAFinal: 0xC4D0
NTDLL32!A_SHAFinal: 0x69400
NTDLL+0x1234: RtlQueryImageMitigationPolicy + 0xf4
ntoskrnl!_DRIVER_OBJECT DeviceObject.
        +0x8    DeviceObject                    : [struct _DEVICE_OBJECT*]
ntoskrnl!_KPROCESS
        +0x0    Header                          : [struct _DISPATCHER_HEADER]
        +0x18   ProfileListHead                 : [struct _LIST_ENTRY]
        +0x28   DirectoryTableBase              : [uint64]
        +0x30   ThreadListHead                  : [struct _LIST_ENTRY]
        ...
"""

symsrv needs to be in the same dir than dbghelp... TODO relax the System32 check