GoLang Argument Parsing and Backtracing

gogo.py

DEBUG = True 

def get_basic_block(ea):
    """get basic blocks of address"""
    f = idaapi.get_func(ea)
    fc = idaapi.FlowChart(f)
    for block in fc:
        if block.start_ea <= ea:
            if block.end_ea > ea:
                return block.start_ea, block.end_ea
    return None, None  


def add_backtrace_comments(arg, index):
    # arg tuple (address, struct variable)
    idc.set_cmt(arg, ("BackTrace Arg:%s" % index), 0)

def add_comments(args):
    # arg tuple (address, struct variable)
    for index, arg in enumerate(args):
        idc.set_cmt(arg[0], ("arg:%s" % index), 0)

def create_arg_comment(arg_data, ea):
    if not arg_data:
        return 
    cmt = []
    for arg in arg_data:
        tt = idc.print_operand(arg[0], arg[1])
        bb = "%s @ 0x%x" % (tt, arg[0])
        cmt.append(bb)
    ff = "(" 
    for cc in cmt:
        ff += cc
        ff += ","
    dd = ff[:-1] # remove trailing , character 
    dd += ")"
    idc.set_cmt(ea, dd, 0)
    
def get_function_args(ea):
    if "call" not in idc.print_insn_mnem(ea):
         return None
    block_start, end = get_basic_block(ea)
    if not block_start:
        return None
    args = [] # offset, stack var
    cur_ea = idc.prev_head(ea)
    while True:
        status, stack_var = get_stack_struct_offset(cur_ea, 0) 
        if not status and stack_var == 0:
            return args   
        if status and stack_var == 0:
            args.insert(0, (cur_ea, stack_var))
            return args
        if status and stack_var != 0:
            args.insert(0, (cur_ea, stack_var))
        cur_ea = idc.prev_head(cur_ea)
        if "call" in idc.print_insn_mnem(cur_ea):
            return None
        if block_start > cur_ea:
            return args 

def get_stack_struct_offset(ea, operand):
    if not ea or ea == idc.BADADDR:
        return False, None     
    ins = ida_ua.insn_t()
    idaapi.decode_insn(ins, ea)
    if ins[operand].type == o_displ:
        op = ins[operand]
        member, stack_offset = ida_frame.get_stkvar(ins, op, op.addr)
        if stack_offset:
            frame_id = idc.get_frame_id(ea)
            if stack_offset > 0:
                oo = stack_offset
                return True, oo
            else:
                oo = stack_offset -16
                return True, oo
                
    # Yuck, "sp" in idc.print_operand(ea, 0)
    # TODO: need to figure out a better approach of code referencing the the stack pointer. 
    # The below code does not identify the code as a stack variale 
    # cur_flag = idc.get_full_flags(ea)
    # idc.is_stkvar0(cur_flag):

    if (ins.itype == idaapi.NN_mov or ins.itype == idaapi.NN_lea) and "sp" in idc.print_operand(ea, 0): 
        return True, 0 
    return False, None 

# todo 
# add stop logic
# track variable types
def go_backtrace(ea):

    # don't backtrace if the second operand is an integer
    ins = ida_ua.insn_t()
    idaapi.decode_insn(ins, ea)
    if ins[1].type == idc.o_imm:
        return ea, 1
    
    block_start, end = get_basic_block(ea)
    last_ref_ea = ea 
    last_ref_op = 1
    trace_str = idc.print_operand(last_ref_ea, last_ref_op)
    last_ref_op = 0
    cur_ea = idc.prev_head(ea)
    while block_start <= cur_ea: 
        if DEBUG:
            print("DEBUG 0:  0x%x, %s %s" % (cur_ea, trace_str, last_ref_op))
        ins = ida_ua.insn_t()
        idaapi.decode_insn(ins, cur_ea)
        if ins.itype == idaapi.NN_call:
            if last_ref_ea == ea:
                return False, None 
            return last_ref_ea, last_ref_op
        # logical operations are not included 
        if ins.itype == idaapi.NN_mov or ins.itype == idaapi.NN_lea:
            if trace_str in idc.print_operand(cur_ea, last_ref_op):
                if DEBUG:
                    print(trace_str, idc.print_operand(last_ref_ea, last_ref_op))
                    print("DEBUG 2:  0x%x, %s %s" % (cur_ea, trace_str, last_ref_op))
                last_ref_ea = cur_ea 
                if last_ref_op == 1:
                    last_ref_op = 0
                else:
                    last_ref_op = 1
                # test type  
                trace_str = idc.print_operand(last_ref_ea, last_ref_op)
                if last_ref_op == 1:
                    last_ref_op = 0
                else:
                    last_ref_op = 1
                if ins[1].type == idc.o_imm:
                    return last_ref_ea, last_ref_op
                if DEBUG:
                    print("DEBUG 3:  0x%x, %s %s" % (cur_ea, trace_str, last_ref_op))
        cur_ea = idc.prev_head(cur_ea)
    
    if last_ref_ea == ea:
        return last_ref_ea, 1 
    
    return last_ref_ea, 1
    

func_addr = here()
args = get_function_args(func_addr)
if args:
    add_comments(args)
    arg_data = [] 
    for index, arg in enumerate(args):
        last_ref_ea, last_ref_op, = go_backtrace(arg[0])
        if DEBUG:
            print(" DEBUG OOO: 0x%x, 0x%x, 0x%x" % (last_ref_ea, last_ref_op, arg[0]))
        if last_ref_ea:
            arg_data.append((last_ref_ea, last_ref_op))
            add_backtrace_comments(last_ref_ea, index)
    create_arg_comment(arg_data, func_addr)

• if source code arg size matches parsed, use source code as reference