import idautils
import idaapi
import idc

print("new ------------------------------------------")

#ea = idc.get_curline()
#print(ea)

# 
def GetWdfVersionBindObject(addr):
    for x in XrefsTo(addr,flags=0):
        cur_addr = x.frm
        cur_asm = GetDisasm(cur_addr)
        if (cur_asm.startswith("call")):
            pass
        else:
            continue

        func_addr = idc.get_func_attr(cur_addr,FUNCATTR_START)
        
        pre_addr = cur_addr
        while True:
            if pre_addr <= func_addr:
                break
            pre_addr = idc.prev_head(pre_addr)
            pre_asm = GetDisasm(pre_addr)
            if (pre_asm.startswith("lea")):
                t = idc.get_operand_type(pre_addr, 0)
                # 寄存器
                if (t == 1):
                    pass
                else:
                    break
            
                data = idc.get_operand_value(pre_addr, 0)
                # r8
                if (data == 8):
                    t = idc.get_operand_type(pre_addr, 1)
                    data = idc.get_operand_value(pre_addr, 1)
                    return data
    return 0


# 从模块中找到对应符号地址
fpWdfVersionBind = idc.get_name_ea(0x140000000, "WdfVersionBind")
print('Address : WdfVersionBind : %#x'%fpWdfVersionBind)
# 根据对应符号地址，找到其第三个参数地址
pObject = GetWdfVersionBindObject(fpWdfVersionBind)
print('Address : Wdf Object : %#x'%pObject)
# 获取的版本号
verBig = idaapi.get_dword(pObject + 0x10)
verMin = idaapi.get_dword(pObject + 0x14)
print("version :", verBig, verMin)


# 根据函数索引取函数名字
def GetNameByID(id):
    return ""


def MakeWdfFunctionInfo(addr):
    for x in XrefsTo(addr,flags=0):
        cur_addr = x.frm
        cur_asm = GetDisasm(cur_addr)
        if (cur_asm.startswith("mov")):
            pass
        else:
            continue
        #print("cur", cur_addr, cur_asm)
        
        pre_addr = idc.prev_head(cur_addr)
        pre_asm = GetDisasm(pre_addr)
        #print("pre", pre_addr, pre_asm)
        if (pre_asm.startswith("imul")):
            # 第 0 个参数类型是1，所以是寄存器
            type0 = idc.get_operand_type(pre_addr, 0)
            if (type0 == 1):
                pass
            else:
                break
            # 寄存器参数索引是 0，是rax
            data = idc.get_operand_value(pre_addr, 0)
            if (data == 0):
                pass
            else:
                break

            # 按理说应该是取操作数1，但是这里1里面没值，所以取的是 2
            # 取出来的就是函数索引
            data = idc.get_operand_value(pre_addr, 2)
            
            # 根据函数索引
            func_name = GetNameByID(data)
            
            fun_addr = idc.get_func_attr(pre_addr,FUNCATTR_START)
            
            if (func_name != ""):
                ida_name.set_name(fun_addr, func_name)


# 寻找所有使用到的地方，并且修正对应函数名，编程索引对应函数名
MakeWdfFunctionInfo(get_qword(pObject + 0x20))