#!/usr/bin/env python

'''

    find_unchecked_alloc.py: Finds calls to a memory allocation routine that doesnt
                             check its return value.
    
    -Cody Pierce
'''

import re
import idaapi

####################################################
#
# Lib calls stuff
#
####################################################

def is_interesting(lib):
    if re.match('.*alloc.*', lib, re.IGNORECASE):
        return True
    
    return False
                   
def get_lib_calls(ea):
    global fh
    refs = []
        
    seg_start = ea
    seg_end = SegEnd(seg_start)

    import_ea = seg_start
    while  import_ea < seg_end:
        import_name = Name(import_ea);
        if len(import_name) > 1:
            xref_start = import_ea
            xref_cur = DfirstB(xref_start)
            while xref_cur != BADADDR:
                if GetFunctionName(xref_cur) in import_name:
                    # We must loop this thunk
                    old_xref_cur = xref_cur
                    old_xref_start = xref_start
                    
                    xref_start = xref_cur
                    xref_cur = RfirstB(xref_start)
                    while xref_cur != BADADDR:
                        if is_interesting(import_name):
                            refs.append((xref_cur, import_name))
                        
                        xref_cur = RnextB(xref_start, xref_cur)
                    
                    xref_cur = old_xref_cur
                    xref_start = old_xref_start
                elif "mov" in GetMnem(xref_cur):
                    reg = GetOpnd(xref_cur, 0)
                    
                    # We must find any references to this indirect call
                    start = xref_cur
                    end = FindFuncEnd(xref_cur)
                    curea = NextHead(start, end)
                    while curea != BADADDR:
                        mnem = GetMnem(curea)
                        if mnem == "call":
                            if GetOpnd(curea, 0) == reg:
                                if is_interesting(import_name):
                                    refs.append((curea, import_name))
                                    
                        curea = NextHead(curea, end)
                else:
                    if is_interesting(import_name):
                        refs.append((xref_cur, import_name))
                        
                xref_cur = DnextB(xref_start, xref_cur)

        import_ea += 4

    return refs
    
# mov     ebx, eax
# cmp     ebx, esi

tests = ["test", "cmp"]
jumps = ["jz", "jl", "jnz"]

refs = get_lib_calls(SegByName(".idata"))
for x in refs:
    ea = x[0]
    name = x[1]
        
    cur_ea = Rfirst(ea)
    retreg = ["eax"]
    tested = False
    jumped = False
    while cur_ea != BADADDR:
        mnem = GetMnem(cur_ea)
        if mnem == "mov":
            op1 = GetOpnd(cur_ea, 0)
            op2 = GetOpnd(cur_ea, 1)
            if op2 in retreg:
                pretreg = retreg[-1]
                retreg.append(GetOpnd(cur_ea, 0))
            elif op1 in retreg:
                retreg.remove(op1)
            elif not len(retreg):
                break
                
        elif mnem in tests:
            op1 = GetOpnd(cur_ea, 0)
            op2 = GetOpnd(cur_ea, 1)
            
            for r in retreg:
                if r in [op1, op2]:
                    tested = True
                    break
        elif mnem in jumps:
            if tested:
                jumped = True
                break
            else:
                dest = idaapi.get_instruction_operand(idaapi.get_current_instruction(), 0).addr
                cur_ea = dest
        elif mnem == "jmp":
            dest = idaapi.get_instruction_operand(idaapi.get_current_instruction(), 0).addr
            cur_ea = dest
            
        cur_ea = Rfirst(cur_ea)
        
    if not jumped:
        print "%x: Did not check its return value" % (ea)