import idaapi
import idautils
import idc
import re

from util import *

def get_func_start(addr):
    func = idaapi.get_func(addr)
    return func.startEA

def get_func_end(addr):
    func = idaapi.get_func(addr)
    return func.endEA

def find_calls(name):
    global functions
    try:
        addrs = [x.frm for x in idautils.XrefsTo(functions['.' + name])]
    except KeyError:
        addrs = []
    return addrs

def previous_ea(ea):
    return idautils.DecodePrecedingInstruction(ea)[0].ip

def next_ea(ea):
    size = idautils.DecodeInstruction(ScreenEA()).size
    return ea+size

def stack_offset(ea):
    ea = int(get_func_start(ea))
    while "sub" not in idc.GetDisasm(ea):
        ea = next_ea(ea)
    return hex2int(idc.GetDisasm(ea).split(',')[1].strip())

def trace_to_input(ea):
    trace = idc.GetDisasm(ea).split(',')[1].strip()
    trace = trace.replace('[', "")
    trace = trace.replace(']', "")
    
    ea = previous_ea(ea)
    instr = idc.GetDisasm(ea).split(',')[0]
    while (trace not in instr) or ('mov' not in instr):
        ea = previous_ea(ea)
        instr = idc.GetDisasm(ea).split(',')[0]

    return ea


def walk_stack(ea, func_args):
    found_args = []
    while len(found_args) < func_args:
        ea = previous_ea(ea)
        if idc.Comment(ea):
            found_args.append((ea, idc.Comment(ea)))

    return found_args

def is_register(reg):
    regs = ['ebp', 'esp', 'eax', 'ebx', 'ecx', 'edx', 'ax', 'bx', 'cx', 'dx', 'ah', 'al', 'bh', 'bl', 'ch', 'cl', 'dh', 'dl']
    for x in regs:
        if x in reg:
            return True
    return False

def trace_arg(ea):
    trace = idc.GetDisasm(ea).split(';')[0].split(',')[1].strip()
    if is_register(trace):
        ea = previous_ea(ea)
        while trace not in idc.GetDisasm(ea):
            ea = previous_ea(ea)
        return ea
    else:
        if trace.isdigit():
            return int(trace)
        else:
            return trace

def hex2int(hexnum):
    return int(hexnum[:-1], 16)

def calculate_buffer_size(ea):
    instr = idc.GetDisasm(ea)
    buf = instr.split(',')[1].strip()
    if "+" in buf:
        try:
            size = hex2int(buf[buf.find('+')+1:len(buf)-1])
        except ValueError:
            old = buf[buf.find('+')+1:len(buf)]
            idc.OpHex(ea, -1)
            instr = idc.GetDisasm(ea)
            buf = instr.split(',')[1].strip()
            size = hex2int(buf[buf.find('-')+1:len(buf)-1])
        return size
        # return stack_offset(ea) - size
    elif "-" in buf:
        size = hex2int(buf[buf.find('-')+1:len(buf)-1])
        return size
        # return stack_offset(ea) + size

BUFFER_SIZE = 0
def stack_calc(function):
    ''' Calculates stack variable sizes for a given function '''
    stack_frame = GetFrame(function) 
    frame_size = GetStrucSize(stack_frame)
    frame_counter = 0
    flag = -1

    stack_vars = {}
    while frame_counter < frame_size:
        stack_var = GetMemberName(stack_frame, frame_counter)
        if stack_var != None:
            if flag != -1:
                size = frame_counter - flag
                if stack_var == " s":
                    size -= 8
                if BUFFER_SIZE < size:
                    stack_vars[current_member] = size
                flag = frame_counter
                current_member = stack_var
            else:
                flag = frame_counter
                current_member = stack_var
            try:
                frame_counter += GetMemberSize(stack_frame, frame_counter)
            except:
                frame_counter += 1
        else:
            frame_counter += 1
    if BUFFER_SIZE < size:
        stack_vars[current_member] = size
    return stack_vars

def test_fgets(ea):
    global bss
    args = walk_stack(ea, 3)
    arg_dict ={}
    for x in args:
        arg_dict[x[1]] = trace_arg(int(x[0]))
    arg_dict['n'] = hex2int(arg_dict['n'])

    try:
        if "offset" in arg_dict['s']:
            section = arg_dict['s'].split(' ')[1]
            arg_dict['s'] = bss[section]        
    except TypeError:
        arg_dict['s'] = calculate_buffer_size(arg_dict['s'])

    if arg_dict['s'] < arg_dict['n']:
        print 'Found vulnerable fgets'
        print 'buf Size: ' + str(arg_dict['s'])
        return True, hex(ea), arg_dict
    else:
        return False, hex(ea), arg_dict

def test_gets(ea):
    return True, hex(ea)

def test_printf(ea):
    args = walk_stack(ea, 1)
    trace = trace_arg(args[0][0])
    arg = idc.GetDisasm(trace).split(',')[1].strip()
    arg_dict = {}
    arg_dict['format'] = arg

    if "esp" in arg:
        return True, hex(ea), arg_dict
    else:
        return False, hex(ea), arg_dict


def test_strcpy(ea):
    args = walk_stack(ea, 2)
    arg_dict = {}
    for x in args:
        arg_dict[x[1]] = x[0]
    arg_dict['dest'] = calculate_buffer_size(trace_arg(arg_dict['dest']))

    arg_dict['src'] = trace_arg(arg_dict['src'])
    if "ebp" in idc.GetDisasm(arg_dict['src']) or "esp" in idc.GetDisasm(arg_dict['src']):
        arg_dict['src'] = -1
    else:
        arg_dict['src'] = trace_to_input(arg_dict['src'])

    if arg_dict['src'] == -1 or "arg" in idc.GetDisasm(arg_dict['src']):
        return True, hex(ea), arg_dict
    else:
        return False, hex(ea), arg_dict

def test_recv(ea):
    args = walk_stack(ea, 4)

    arg_dict = {}
    for x in args:
        arg_dict[x[1]] = trace_arg(x[0])
    buf = arg_dict['buf']

    try:
        buf = stack_calc(ea)['buf']
    except KeyError:
        pass

    arg_dict['buf'] = buf

    if type(arg_dict['n']) == str:
        arg_dict['n'] = hex2int(arg_dict['n'])

    if arg_dict['n'] > buf:
        return True, hex(ea), arg_dict
    else:
        return False, hex(ea), arg_dict

def test_read(ea):
    args = walk_stack(ea, 3)
    arg_dict = {}
    for x in args:
        arg_dict[x[1]] = trace_arg(x[0])
    arg_dict['nbytes'] = hex2int(arg_dict['nbytes'])
    arg_dict['buf'] = calculate_buffer_size(arg_dict['buf'])

    if arg_dict['buf'] < arg_dict['nbytes']:
        return True, hex(ea), arg_dict
    else:
        return False, hex(ea), arg_dict


### Populate BSS <Hail Chris Eagle>
start = idaapi.get_segm_by_name(".bss").startEA
end = idaapi.get_segm_by_name(".bss").endEA
item = idc.NextHead(start - 1, end)
bss = {}
while item != BADADDR:
   next = idc.NextHead(item, end)
   if next != BADADDR:
        bss[idc.Name(item)] = next - item
   else:
        bss[idc.Name(item)] = end - item
   item = next

### Populate Functions
functions = {}
for x in idautils.Functions():
    functions[idc.GetFunctionName(x)] = x


print "gets", find_calls('gets')
for x in find_calls('gets'):
    res = test_gets(x)
    print '\t', res
    if res[0]:  idc.SetColor(int(res[1], 16), CIC_ITEM, 0x0000ff)


print "fgets", find_calls('fgets')
for x in find_calls('fgets'):
    res = test_fgets(x)
    print '\t', res
    if res[0]:  idc.SetColor(int(res[1], 16), CIC_ITEM, 0x0000ff)


print "strcpy", find_calls('strcpy')
for x in find_calls('strcpy'):
    res = test_strcpy(x)
    print '\t', res
    if res[0]:  idc.SetColor(int(res[1], 16), CIC_ITEM, 0x0000ff)


print "recv", find_calls('recv')
for x in find_calls('recv'):
    res = test_recv(x)
    print '\t', res
    if res[0]:  idc.SetColor(int(res[1], 16), CIC_ITEM, 0x0000ff)


print "printf", find_calls('printf')
for x in find_calls('printf'):
    res = test_printf(x)
    print '\t', res
    if res[0]:  idc.SetColor(int(res[1], 16), CIC_ITEM, 0x0000ff)


print "read", find_calls('read')
for x in find_calls('read'):
    res = test_read(x)
    print '\t', res
    if res[0]:  idc.SetColor(int(res[1], 16), CIC_ITEM, 0x0000ff)


