#coding: utf-8
import ida_dbg
import idc
import idaapi
import funcutils

class Advapi32:
    def __init__(self):
        self.is_64bit = True if idaapi.BADADDR == 0xffffffffffffffff else False
        self.total_func = {
            "RegCreateKeyA": self.RegCreateKeyA,
            "RegCreateKeyW": self.RegCreateKeyW,
        }

    def RegCreateKeyA(self, is_return = False):
        if is_return == True:
            phkResult = funcutils.get_func_args(2, self.is_64bit)
            if phkResult != 0:
                result = 0
                if self.is_64bit == True:
                    result = idc.get_qword(phkResult)
                else:
                    result = idc.get_wide_dword(phkResult)
            return f"{hex(funcutils.get_result(self.is_64bit))} -> LSTATUS, hkResult = {hex(result)}"
        
        hKey = funcutils.get_func_args(1, self.is_64bit)
        lpSubKey = funcutils.get_func_args(2, self.is_64bit)
        phkResult = funcutils.get_func_args(3, self.is_64bit)
        if hKey == 0x80000000:
            _debug_info = f"hKey={hex(hKey)} -> HKEY_CLASSES_ROOT"
        elif hKey == 0x80000001:
            _debug_info = f"hKey={hex(hKey)} -> HKEY_CURRENT_USER"
        elif hKey == 0x80000002:
            _debug_info = f"hKey={hex(hKey)} -> HKEY_LOCAL_MACHINE"
        elif hKey == 0x80000003:
            _debug_info = f"hKey={hex(hKey)} -> HKEY_USERS"
        elif hKey == 0x80000004:
            _debug_info = f"hKey={hex(hKey)} -> HKEY_PERFORMANCE_DATA"
        else:
            _debug_info = f"hKey={hex(hKey)}"
        
        _debug_info += f", lpSubKey={idc.get_strlit_contents(lpSubKey).decode('utf-8')}, "
        _debug_info += f", phkResult={hex(phkResult)}) = "
        return _debug_info

    def RegCreateKeyW(self, is_return = False):
        if is_return == True:
            phkResult = funcutils.get_func_args(2, self.is_64bit)
            if phkResult != 0:
                result = 0
                if self.is_64bit == True:
                    result = idc.get_qword(phkResult)
                else:
                    result = idc.get_wide_dword(phkResult)
            return f"{hex(funcutils.get_result(self.is_64bit))} -> LSTATUS, hkResult = {hex(result)}"
        
        hKey = funcutils.get_func_args(1, self.is_64bit)
        lpSubKey = funcutils.get_func_args(2, self.is_64bit)
        phkResult = funcutils.get_func_args(3, self.is_64bit)
        if hKey == 0x80000000:
            _debug_info = f"hKey={hex(hKey)} -> HKEY_CLASSES_ROOT"
        elif hKey == 0x80000001:
            _debug_info = f"hKey={hex(hKey)} -> HKEY_CURRENT_USER"
        elif hKey == 0x80000002:
            _debug_info = f"hKey={hex(hKey)} -> HKEY_LOCAL_MACHINE"
        elif hKey == 0x80000003:
            _debug_info = f"hKey={hex(hKey)} -> HKEY_USERS"
        elif hKey == 0x80000004:
            _debug_info = f"hKey={hex(hKey)} -> HKEY_PERFORMANCE_DATA"
        else:
            _debug_info = f"hKey={hex(hKey)}"
        
        _debug_info += f", lpSubKey={idc.get_strlit_contents(ea = lpSubKey, strtype=idc.STRTYPE_C16).decode('utf-8')}, "
        _debug_info += f", phkResult={hex(phkResult)}) = "
        return _debug_info

    def RegCreateKeyExA(self, is_return = False):
        pass
    
    def RegCreateKeyExW(self, is_return = False):
        pass

    def RegSetValueA(self, is_return = False):
        pass

    def RegSetValueW(self, is_return = False):
        pass