IDAPython snippets

_IDAPython-Snippets.md

This gist contains lots of my written IDAPython snippets.

fix_arm_xref.py

# partially dirty fix for non-adjacent MOVW + MOVT.W combination.
from idaapi import *
cur = 0x89cc
ea = 0x8a0e
while True:
    if print_insn_mnem(ea) == "MOVW":
        decode_insn(ea)
        high = cmd[1].value
        high_reg = cmd[0].reg
        rev = ea
        for i in range(6):
            rev = PrevHead(rev)
            if print_insn_mnem(rev) == "MOVT.W":
                decode_insn(rev)
                if cmd[0].reg == high_reg:
                    target = (high) | (cmd[1].value << 16)
                    print "Got two phase mov at %x, value %x" % (ea, target)
                    if target < 0x80000:
                        cmd.add_dref(target, 1, dr_O)
                break
    elif print_insn_mnem(ea) == "MOVT.W":
        decode_insn(ea)
        high = cmd[1].value
        high_reg = cmd[0].reg
        rev = ea
        for i in range(6):
            rev = PrevHead(rev)
            if print_insn_mnem(rev) == "MOVW":
                decode_insn(rev)
                if cmd[0].reg == high_reg:
                    target = (high << 16) | cmd[1].value
                    print "Got two phase mov at %x, value %x" % (ea, target)
                    if target < 0x80000:
                        cmd.add_dref(target, 1, dr_O)
                break
    
 
    ea = NextHead(ea)
    if ea == BADADDR or ea > 0x50000:
        break

flare-emu_memhook.py

import sys
sys.path.append(r'D:\Workspaces\UtilWorkspace\Reverse\IDAPlugins\flare-emu')

import flare_emu    
import unicorn as UC

def mem_hook(unicornObject, accessType, memAccessAddress, memAccessSize, memValue, userData):
    #if accessType == UC.UC_MEM_READ:
    #    print("Read: ", hex(memAccessAddress), memAccessSize, hex(memValue))
    if accessType == UC.UC_MEM_WRITE:
        #print("Write: ", hex(memAccessAddress), memAccessSize, hex(memValue))
        if memAccessSize == 1:
            idc.PatchByte(memAccessAddress, memValue)
        elif memAccessSize == 2:
            idc.PatchWord(memAccessAddress, memValue)
        elif memAccessSize == 4:
            idc.PatchDword(memAccessAddress, memValue)

def getInitEh():
    eh = flare_emu.EmuHelper()
    eh.addApiHook('LoadLibraryA', 'GetProcessHeap')
    eh.addApiHook('GetProcAddress', 'GetProcessHeap')
    eh.addApiHook('ClearLibName', 'GetProcessHeap')
    eh.addApiHook('SanitizeString', 'GetProcessHeap')
    eh.addApiHook('j_RtlAllocateHeap', lambda eh, address, argv, funcName, userData: eh.uc.reg_write(eh.regs["ret"], 0x600000))
    eh.emulateRange(0x00407EB9, endAddr=0x407ef6, skipCalls=False, memAccessHook=mem_hook)
    return eh

eh = getInitEh()

# -----------------------------------------------------

def call_hook(address, argv, funcName, userData):
    print('{}:call fun:{}'.format(hex(address),funcName))
    print("argv:", hex(argv[0]), argv[1:])
    eh = userData["EmuHelper"]
    if funcName == 'GetProcAddress':
        curEdi = eh.uc.reg_read(eh.regs['edi'])
        curProc = eh.getEmuString(argv[1])
        curProc = curProc.decode('latin-1')
        print("GetProc to " + hex(curEdi) + ": "+ curProc)
        MakeName(curEdi, "j_" + curProc)

eh.emulateRange(0x407ef6, endAddr=0x407efb, skipCalls=False, callHook=call_hook)
eh.emulateRange(0x40289B, endAddr=0x402911, skipCalls=False, memAccessHook=mem_hook, callHook=call_hook)

# -----------------------------------------------------

def decstr_call_hook(address, argv, funcName, userData):
    #print('{}:call fun:{}'.format(hex(address),funcName))
    #print("argv:", hex(argv[0]), argv[1:])
    eh = userData["EmuHelper"]
    if funcName == 'DecryptStringInplace':
        userData['decLoc'] = argv[0]

d = {}
for x in XrefsTo(get_name_ea(BADADDR, 'DecryptStringInplace')):
    ea1 = PrevHead(x.frm)
    ea2 = PrevHead(ea1)
    ea3 = PrevHead(ea2)
    start = None
    if Byte(ea1) == 0x68 and Byte(ea2) == 0xFF:
        start = ea2
    elif Byte(ea1) == 0x53 and Byte(ea2) == 0xff and Byte(ea3) == 0x8D:
        start = ea3
    else:
        continue
    dd = GetDisasm(start)
    if dd in d:
        continue
    d[dd] = 1
    print("Handling %x" % start)
    eh.emulateRange(start, endAddr=NextHead(x.frm), skipCalls=False, memAccessHook=mem_hook, callHook=decstr_call_hook)
    PatchDword(eh.hookData['decLoc'] - 4, 0)

get_all_user_define_names.py

from idaapi import *
from idc import *
import idautils

def getUserDefinedName():
    for ea, name in idautils.Names():
        if has_user_name(getFlags(ea)):
            yield (ea,name)

def main():
    for ea, name in getUserDefinedName():
        print "ea: %X name:%s" % (ea, name)

main()

get_basic_block_end.py

# This snippet also supports bb that is not belonged to functions
bb_end = bb_addr
while True:
    decode_insn(bb_end)
    if is_basic_block_end(bb_end):
        break
    bb_end = NextHead(bb_end)

handle_str_obfs.py

buf = '''
byte_10A15 ^= 0xD3u;
    byte_10A16 ^= 0x41u;
    byte_10A17 ^= 0x31u;
    byte_10A18 ^= 0xDBu;
    byte_10A19 ^= 0x94u;
    byte_10A1A ^= 0xDu;
    byte_10A1B ^= 0x12u;
    byte_10A1C ^= 0x18u;
    byte_10A1D ^= 0x68u;
    byte_10A1E ^= 0xB0u;
    byte_10A20 ^= 0x3Cu;
    byte_10A21 ^= 0x3Fu;
    byte_10A22 ^= 0x6Au;
    byte_10A23 ^= 0x89u;
    byte_10A24 ^= 0x22u;
    byte_10A25 ^= 0x83u;
    byte_10A26 ^= 0xA1u;
    byte_10A27 ^= 0xD0u;
    byte_10A28 ^= 0xDCu;
    byte_10A29 ^= 0xB9u;
    byte_10A2A ^= 0xBFu;
    byte_10A2B ^= 0x44u;
    byte_10A2C ^= 0xC6u;
    byte_10A2D ^= 0x79u;
    byte_10A2E ^= 0xE6u;
    byte_10A2F ^= 0xB9u;
    byte_10A30 ^= 0x52u;
    byte_10A31 ^= 0x6Du;
    byte_10A32 ^= 0x78u;
    byte_10A33 ^= 0x29u;
    byte_10A34 ^= 0xB3u;
    byte_10A35 ^= 0x31u;
    byte_10A36 ^= 0x65u;
    byte_10A37 ^= 0xE1u;
    byte_10A38 ^= 0x7Cu;
    byte_10A39 ^= 0x61u;
    byte_10A3A ^= 0xE2u;
    byte_10A3B ^= 0x48u;
    byte_10A3C ^= 0x62u;
    byte_10A3D ^= 0xB2u;
    byte_10A3E ^= 0xB1u;
    byte_10A3F ^= 0x56u;
    byte_10A40 ^= 0xE1u;
    byte_10A41 ^= 9u;
    byte_10A42 ^= 0xE2u;
    byte_10A43 ^= 0x6Cu;
    byte_10A44 ^= 0x32u;
    byte_10A45 ^= 0xB5u;
    byte_10A46 ^= 0x85u;
    byte_10A47 ^= 0x79u;
    byte_10A48 ^= 0x67u;
    byte_10A49 ^= 0x14u;
    byte_10A50 ^= 0x75u;
    byte_10A51 = ~byte_10A51;
    byte_10A52 ^= 0xA0u;
    byte_10A53 ^= 0xE6u;
    byte_10A54 ^= 0x7Fu;
    byte_10A55 ^= 0x9Fu;
    byte_10A56 ^= 0x7Cu;
    byte_10A57 ^= 0x84u;
    byte_10A58 ^= 9u;
    byte_10A59 ^= 0xEu;
    byte_10A5B ^= 0xC2u;
    byte_10A5C ^= 0x8Cu;
    byte_10A5D ^= 0x48u;
    byte_10A5E ^= 0x2Bu;
    byte_10A5F ^= 0xF9u;
    byte_10A60 ^= 0xF7u;
    byte_10A61 ^= 0x4Du;
    byte_10A62 ^= 0x4Fu;
    byte_10A63 ^= 0xDCu;
    byte_10A64 ^= 0xC8u;
    byte_10A65 ^= 0x7Au;
    byte_10A66 ^= 0x47u;
    byte_10A67 ^= 0x9Fu;
    byte_10A68 ^= 0x59u;
    byte_10A69 ^= 0x9Bu;
    byte_10A6A ^= 0x1Eu;
    byte_10A6B ^= 0xF7u;
'''
def get_name_type(name):
	if name.startswith("byte"):
		return "Byte"
	elif name.startswith("qword"):
		return "Qword"
	else:
		return None
def handle_buf(buf):
    for line in buf.splitlines():
    	line = line.strip().rstrip(";")
    	print "Processing line: %s" % line
    	if "^=" in line:
    		name, _, num = line.partition("^=")
    		addr = get_name_ea(BADADDR,name.strip())
    		num = num.strip().rstrip("L").rstrip("u")
    		imm = int(num, 16) if num.startswith("0x") else int(num)
    		nametype = get_name_type(name)
    		eval("Patch"+nametype+"(addr, "+ nametype + "(addr) ^ imm)")
    	elif " = ~" in line:
    		print "Handling not"
    		name, _, _ = line.partition("=")
    		name.strip()
    		addr = get_name_ea(BADADDR,name.strip())
    		nametype = get_name_type(name)
    		print "Patch"+nametype+"(addr, ~"+ nametype + "(addr))"
    		eval("Patch"+nametype+"(addr, ~"+ nametype + "(addr))")
 
 handle_buf(buf)

ida_debug_daemon.py

import sys
if not r"D:\ProjWorkspace\example_mod" in sys.path:
    sys.path.append(r"D:\ProjWorkspace\example_mod")
    sys.path.append(r"C:\Program Files\JetBrains\PyCharm 2019.2.1\debug-eggs\pydevd-pycharm.egg")

import idaapi

def hotkey_pressed():
    idaapi.require("example_mod")
    import pydevd
    pydevd.settrace('localhost', port=11100, stdoutToServer=True, stderrToServer=True, suspend=False)
    example_mod.main()
    pydevd.stoptrace()


hotkey_ctx = idaapi.add_hotkey("Ctrl-Shift-A", hotkey_pressed)

instruction_pattern_replace.py

import re
from keystone import *

ks = Ks(KS_ARCH_X86, KS_MODE_64)

ea = 0x1000007a0
ea_end = 0x100013f80

while ea < ea_end:
    ea1 = ea
    ea2 = next_head(ea1)
    ea3 = next_head(ea2)
    ea = ea2
    
    m = re.match(r'mov\s*(r[a-z0-9]+?), rsp', GetDisasm(ea1))
    if not m:
        continue
    reg1 = m.group(1)
    
    m = re.match(r'add\s*(r[a-z0-9]+?), (.*?)h', GetDisasm(ea2))
    if not m:
        continue
    reg2 = m.group(1)
    spd = m.group(2)
            
    m = re.match(r'mov\s*rsp, (r[a-z0-9]+?)$', GetDisasm(ea3))
    if not m:
        continue
    reg3 = m.group(1)
    if reg1 != reg2 or reg2 != reg3:
        continue
        
    retbuf, retinsnum = ks.asm(f"""
    add rsp, {spd}h; mov {reg1}, rsp;
    """)
    print(retbuf, retinsnum)
    #assert retinsnum == 2
    ret = b''.join((c.to_bytes(1, 'little') for c in retbuf))
    ea = next_head(ea3)
    buflen = ea - ea1
    buf = ret.rjust(buflen, b'\x90')
    idaapi.patch_bytes(ea1, buf)
    print("patched %x" % ea1)

logfunc_rename.py

# Rename function based on call to logger functions such as assert("assert failed", "function name", line_num, ...)
from collections import defaultdict

class CallTargetFinder(idaapi.ctree_visitor_t):
    def __init__(self, target, cb):
        idaapi.ctree_visitor_t.__init__(self, idaapi.CV_FAST)
        self.target = target
        self.cb = cb
    
    # Hex-Rays will call this function for every expression within the first
    # line of code extracted from the case statement bodies.
    def visit_expr(self,i):
        # Look for call expressions...
        if i.op == idaapi.cot_call:
            # ... where the target is a known location in the database.
            if i.x.op == idaapi.cot_obj and i.x.obj_ea == self.target:
                self.cb(i)

        return 0


def process_func(ea, target, argloc):
    f = idaapi.get_func(ea)
    if f is None:
        print "Couldn't get func_t for hexdsp?"
        return

    # Decompile the function.
    cfunc = idaapi.decompile(f)
    if cfunc is None:
        print "Failed to decompile!"
        return

    def handler(i):
        nameexpr = i.a[argloc]
        if nameexpr.op == idaapi.cot_str:
            newname = nameexpr.string
        elif nameexpr.op == idaapi.cot_obj:
            newname = get_strlit_contents(nameexpr.obj_ea)
        else:
            print "Warning: Failed to get arg %d as string!" % argloc
            return
            
        print "Renaming function at %x to %s" % (f.startEA, newname)
        idaapi.set_name(f.startEA, newname, SN_AUTO | SN_NOCHECK | SN_NOWARN)

    # Apply the visitors above to the decompilation listing.
    se = CallTargetFinder(target, handler)
    se.apply_to(cfunc.body, None)


def process_xrefs(logfun, argloc):
    # argloc starts from 0
    for xref in XrefsTo(logfun, 0):
        print "Processing xref: %s(0x%x)" % (GetFunctionName(xref.frm), xref.frm)    
        ref_type = ""
        try:
            process_func(xref.frm, logfun, argloc)
        except:
            pass

process_xrefs(0x2eb24, 1)
process_xrefs(0x2d250, 2)

make_all_functions.py

# This version supports the function chunk, which will be much more robust than func.endEA
def make_all_func(start,end):
    cur = start
    while True:
        print("Making %x" % cur)
        idaapi.add_func(cur)
        endchunk = idc.get_fchunk_attr(cur, FUNCATTR_END)
        cur = endchunk if endchunk != BADADDR else idc.next_addr(cur)
        if cur > end:
            break
            
make_all_func(0x60010000, 0x60cd2690)

make_func_sig.py

# get hex pattern of a function (excluding relocations)

func = get_func(ScreenEA())
sig = ""
config = Config()
"""
type config: Config
type func: idc.func_t
"""
logger = logging.getLogger("idb2pat:make_func_sig")

if func.end_ea - func.start_ea < config.min_func_length:
    logger.debug("Function is too short")
    raise FuncTooShortException()

ea = func.start_ea
publics = []  # type: idc.ea_t
refs = {}  # type: dict(idc.ea_t, idc.ea_t)
variable_bytes = set([])  # type: set of idc.ea_t

while ea != BADADDR and ea < func.end_ea:
    logger.debug("ea: %s", hex(ea))

    name = get_name(ea)
    if name is not None and name != "":
        logger.debug("has a name")
        publics.append(ea)

    ref = get_first_dref_from(ea)
    if ref != BADADDR:
        # data ref
        logger.debug("has data ref")
        ref_loc = find_ref_loc(config, ea, ref)
        if ref_loc != BADADDR:
            logger.debug("  ref loc: %s", hex(ref_loc))
            for i in zrange(config.pointer_size):
                logger.debug("    variable %s", hex(ref_loc + i))
                variable_bytes.add(ref_loc + i)
            refs[ref_loc] = ref

        # not sure why we only care about two...
        # only two possible operands?
        ref = get_next_dref_from(ea, ref)
        if ref != BADADDR:
            logger.debug("has data ref2")
            ref_loc = find_ref_loc(config, ea, ref)
            if ref_loc != BADADDR:
                logger.debug("  ref loc: %s", hex(ref_loc))
                for i in zrange(config.pointer_size):
                    logger.debug("    variable %s", hex(ref_loc + i))
                    variable_bytes.add(ref_loc + i)
                refs[ref_loc] = ref
    else:
        # code ref
        ref = get_first_fcref_from(ea)
        if ref != BADADDR:
            logger.debug("has code ref")
            if ref < func.start_ea or ref >= func.end_ea:
                # code ref is outside function
                ref_loc = find_ref_loc(config, ea, ref)
                if BADADDR != ref_loc:
                    logger.debug("  ref loc: %s", hex(ref_loc))
                    for i in zrange(config.pointer_size):
                        logger.debug("    variable %s", hex(ref_loc + i))
                        variable_bytes.add(ref_loc + i)
                    refs[ref_loc] = ref

    ea = next_not_tail(ea)

for ea in zrange(func.start_ea, func.end_ea):
    if ea in variable_bytes:
        sig += "?? "
    else:
        sig += "%02X " % (get_byte(ea))

print(sig)

mark_all_offset.py

RANGE_START = 0x6006f544
RANGE_END = 0x60cd2690
ADDR_START = 0x60010000
ADDR_END = 0x60cd2690
for ea in range(RANGE_START, RANGE_END, 4):
    word = idaapi.get_dword(ea)
    if word >= ADDR_START and word < ADDR_END:
        idaapi.del_items(ea)
        idaapi.op_plain_offset(ea, 0, 0)

mark_as_decompiled.py

# from https://github.com/gaasedelen/lighthouse/issues/27
def mark_as_decompiled(function_address):
    function = idaapi.get_func(function_address)
    if not function:
        return False
    function.color = 0xEEFFF0
    return True

ea = 0x4280010C0 - 1
while True:
    ea = NextFunction(ea)
    if ea > 0x428006F40:
        break
    print ea,mark_as_decompiled(ea)

more_compilcated_ctree_example.py

from collections import defaultdict

class FindSEPCall(idaapi.ctree_visitor_t):
    def __init__(self, sepworker):
        idaapi.ctree_visitor_t.__init__(self, idaapi.CV_FAST)
        self.sepworker = sepworker
        self.call = []

    # Hex-Rays will call this function for every expression within the first
    # line of code extracted from the case statement bodies.
    def visit_expr(self,i):
        # Look for call expressions...
        if i.op == idaapi.cot_call and i.x.op == idaapi.cot_obj:
            # ... where the target is a known location in the database.
            if len(i.a) >= 1:
                a0 = i.a[0]
                if a0.op == idaapi.cot_obj and a0.obj_ea == self.sepworker:
                    self.call += [i.x.obj_ea]
        return 0


class FindSEPMsgID(idaapi.ctree_visitor_t):
    def __init__(self, argivar):
        idaapi.ctree_visitor_t.__init__(self, idaapi.CV_FAST)
        self.tracing = [argivar]
        self.foundid = -1

    # Hex-Rays will call this function for every expression within the first
    # line of code extracted from the case statement bodies.
    def visit_expr(self,i):
        # Look for call expressions...
        if i.op == idaapi.cot_asg:
            if i.y.op == idaapi.cot_var and i.y.v.idx in self.tracing:
                self.tracing += [i.x.v.idx]
        if i.op == idaapi.cot_call:
            if i.x.op == idaapi.cot_var and i.x.v.idx in self.tracing:
                assert(i.a  [1].op == idaapi.cot_num)
                self.foundid = i.a[1].n

        return 0


def get_callsep_funcs(sepworker):
    ret = []
    for xref in XrefsTo(sepworker, 0):
        print "Finding target in xref: %s(0x%x)" % (GetFunctionName(xref.frm), xref.frm)    
        cfunc = idaapi.decompile(xref.frm)
        if not cfunc:
            continue
        finder = FindSEPCall(sepworker)
        finder.apply_to(cfunc.body, None)
        ret += finder.call
    return ret

def process_callsep(callsep):
    cfunc = idaapi.decompile(callsep)
    sepvar = None
    for i,c in enumerate(cfunc.get_lvars()):
        if c.is_arg_var:
            sepvar = i
            break
    finder = FindSEPMsgID(sepvar)
    finder.apply_to(cfunc.body, None)
    if finder.foundid._value != -1:
        newname = "callsep_%d" % finder.foundid._value
        print "Renamed %x to %s" % (callsep, newname)
        idaapi.set_name(callsep, newname, SN_NOCHECK | SN_NOWARN | SN_AUTO)
    else:
        print "Failed to find msgid for %x" % callsep
       

fun = get_callsep_funcs(0xFFFFFFF0067CD14C)
print fun
for c in fun:
    process_callsep(c)