IDAPython Utils

IDAPythonUtils.py

# The MIT License (MIT)
# Copyright (C) 2018 5kyc1ad
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

from idaapi import *

REGISTER_NAMES = ["eax", "ebx", "ecx", "edx", "edi", "esi"]

ARGTYPE_OFFSET = 0
ARGTYPE_STACK_OFFSET = 1
ARGTYPE_LONG = 2

def IsIDAHexString(targetString):
    try:
        int(targetString.rstrip('h'), 16)
    except:
        return False
    return True

def GetStackVars():
    definedValues = {"Stack": {}, "Global": {}}

    # EA : Effective Address
    ea = here()
    funcStart = GetFunctionAttr(ea, FUNCATTR_START)
    funcEnd =  GetFunctionAttr(ea, FUNCATTR_END)

    currentEA = funcStart
    while True:
        disasm = GetDisasm(currentEA)
        
        if disasm.startswith("mov "):
            lValue = GetOpnd(currentEA, 0)
            rValue = GetOpnd(currentEA, 1)

            if lValue.startswith("[esp") or lValue.startswith("[ebp"):
                lValue = GetOperandValue(currentEA, 0)
                if IsIDAHexString(rValue):
                    definedValues["Stack"][lValue] = int(rValue.rstrip('h'), 16)
                else:
                    definedValues["Stack"][lValue] = definedValues["Global"][rValue] if rValue in definedValues["Global"].keys() else "Unknwon:"+rValue
            else:
                if IsIDAHexString(rValue):
                    definedValues["Global"][lValue] = int(rValue.rstrip('h'), 16)
                else:
                    definedValues["Global"][lValue] = definedValues["Global"][rValue] if rValue in definedValues["Global"].keys() else "Unknwon:"+rValue

        currentEA = NextHead(currentEA)
        if currentEA > funcEnd:
            break
    return definedValues

def GetStrFromAddr(offset, func = None):
    result = ""
    while True:
        value = Byte(offset)
        if func is not None:
            value = func(value)
        if value == 0:
            break
        result += chr(value)
        offset += 1
    return result

def GetStackVarStr(definedVars, offset, func = None):
    #print json.dumps(definedVars, indent = 4, sort_keys = 4)
    #print offset
    result = ""
    while True:
        if offset in definedVars["Stack"].keys():
            value = definedVars["Stack"][offset]
            if func is not None:
                value = func(value)
            if value == 0:
                break
            result += chr(value)
            offset += 1
        else:
            break
    return result

# Return : (ArgType, ArgValue)
def GetArgumentAddr(callAddr, argNum):
    if not GetDisasm(callAddr).startswith("call "):
        raise ValueError("Not a valid call address : 0x%08x" % (callAddr))

    prevAddr = callAddr
    prevCount = 0
    while True:
        prevAddr = PrevHead(prevAddr)
        disasm = GetDisasm(prevAddr)
        if disasm.startswith("push "):
            prevCount += 1
        if disasm.startswith("push ") and prevCount == argNum:
            arg = GetOpnd(prevAddr, 0)
            if type(arg) == str and arg.startswith("offset "):
                return (ARGTYPE_OFFSET, GetOperandValue(prevAddr, 0))
            elif type(arg) == str and arg in REGISTER_NAMES:
                while True:
                    prevAddr = PrevHead(prevAddr)
                    if GetDisasm(prevAddr).startswith("lea ") and GetOpnd(prevAddr, 0) == arg:
                        return (ARGTYPE_STACK_OFFSET, GetOperandValue(prevAddr, 1))
            elif type(arg) == int:
                return (ARGTYPE_LONG, arg)
            else:
                raise Exception("Unknown arg (type: , value: , dir: )" % (type(arg), str(arg), dir(arg)))
            


def RenameGetProcAddrVars(getProcAddr):
    definedVars = GetStackVars()
    ea = here()
    print " =================================== "
    print "[*] Current Address :", hex(ea)
    funcStart = GetFunctionAttr(ea, FUNCATTR_START)
    print "[*] Func Start :", hex(funcStart)
    funcEnd =  GetFunctionAttr(ea, FUNCATTR_END)
    print "[*] Func End :", hex(funcEnd)
    print " =================================== "

    # Get xrefs of 'GetProcAddress' in Current Function
    for xref in filter(lambda xref : funcStart < xref.frm < funcEnd, list(XrefsTo(getProcAddr, flags = 0))):
        instAddr = xref.frm
        if not GetDisasm(instAddr).startswith("call "):
            continue

        argType, argValue = GetArgumentAddr(instAddr, 2)
        if argType == ARGTYPE_OFFSET:
            funcName = GetStrFromAddr(argValue)
        else:
            funcName = GetStackVarStr(definedVars, argValue, lambda x: x ^ 0x08)

        while True:
            instAddr = NextHead(instAddr)
            if GetDisasm(instAddr).startswith("mov ") and GetOpnd(instAddr, 1) == "eax":
                break
        
        print "[+] {0} = GetProcAddress(\"{1}\")".format(hex(GetOperandValue(instAddr, 0)), funcName)
        nameCount = 0
        while True:
            res = set_name(GetOperandValue(instAddr, 0), "p_{0}_{1}".format(funcName, nameCount) if nameCount > 0 else "p_{0}".format(funcName))
            if res == True:
                break
            nameCount += 1
    print " =================================== "