'''
    Author: Alexander Hanel 
    Date: 20171007
    Purpose: calculates a structure members and size based off of offsets into an address 
'''
import idautils 
import idaapi

DISTANCE = 0x8

def get_reg_type(register):
    if register in ['al', 'ah', 'ax', 'eax', 'rax']:
        return 'accumulator' 
    if register in ['bl', 'bh', 'bx', 'ebx', 'rbx']:
        return 'base'
    if register in ['cl', 'ch', 'cx', 'ecx', 'rcx']:
        return 'counter'
    if register in ['dl', 'dh', 'dx', 'edx', 'rdx']:
        return 'extend'
    if register in ['si', 'esi', 'rsi']:
        return 'source'
    if register in ['di', 'edi', 'rdi']:
        return 'dest'
    if register in ['sp', 'esp', 'rbp']:
        return 'stack'
    if register in ['bp', 'ebp', 'rbp']:
        return 'base'
    if register in ['ip', 'eip', 'rip']:
        return 'instru' 
    return None

def get_arch_bit():
    info = idaapi.get_inf_structure()
    if info.is_64bit():
        bits = 64
    elif info.is_32bit():
        bits = 32
    else:
        bits = 16
    return bits
    
def get_to_xrefs(ea):
    xref_set = set([])
    for xref in idautils.XrefsTo(ea, 1):
        xref_set.add(xref.frm) 
    return xref_set

def read_offset(addr, op):
    idaapi.decode_insn(addr)
    if op == 1:
        offset = idaapi.cmd.Op1.addr
    else:
        offset = idaapi.cmd.Op2.addr
    return offset

def get_reg(addr, op):
    """
        get register from displacement 
        example: mov [edx+12Ch], eax
    """
    idaapi.decode_insn(addr)
    bit = get_arch_bit()
    if bit == 16:
        if op == 1:
            return idaapi.get_reg_name(idaapi.cmd.Op1.reg, 2)
        else:
            return idaapi.get_reg_name(idaapi.cmd.Op2.reg, 2)
    if bit == 32:
        if op == 1:
            return idaapi.get_reg_name(idaapi.cmd.Op1.reg, 4)
        else:
            return idaapi.get_reg_name(idaapi.cmd.Op2.reg, 4)
    if bit == 64:
        if op == 1:
            return idaapi.get_reg_name(idaapi.cmd.Op1.reg, 8)
        else:
            return idaapi.get_reg_name(idaapi.cmd.Op2.reg, 8)
    return None

def add_struct(size_offset):
    sid = idc.AddStrucEx(-1, 'name_me', 0)
    for count, offset in enumerate(size_offset):
        struct_name = "field_%x" % (count)
        if count == len(size_offset) - 1:
            # unknown size
            idc.AddStrucMember(sid, struct_name, offset, FF_BYTE, -1, 1)
        else:
            size =  size_offset[count +1] - offset
            idc.AddStrucMember(sid, struct_name, offset, FF_BYTE, -1, size)
        

def forward_trace(addr, og_reg):
    global DISTANCE
    og_reg_type = get_reg_type(og_reg)
    # forward trace n (DISTANCE) from xref to struct offset 
    for x in xrange(0, DISTANCE):
        idaapi.decode_insn(addr)
        # TODO: add mov and calcualte lea check 
        # check for write to struct value
        reg_type_op1 = get_reg_type(get_reg(addr, 1))
        if idaapi.cmd.Op1.type == idaapi.o_displ:
            if og_reg_type == reg_type_op1:
                temp_offset = read_offset(addr, 1) 
                return (True, temp_offset, addr)
        # check for read to struct value
        reg_type_op2 = get_reg_type(get_reg(addr, 2))
        if idaapi.cmd.Op2.type == idaapi.o_displ:
            if og_reg_type == reg_type_op2:
                temp_offset = read_offset(addr, 2) 
                return (True, temp_offset, addr)
        # check for 0 offset
        if idaapi.cmd.Op1.type == idaapi.o_phrase:
            if og_reg_type == reg_type_op1:
                return (True, 0,addr)
        if idaapi.cmd.Op2.type == idaapi.o_phrase:
             if og_reg_type == reg_type_op2:
                return (True, 0,addr)
        addr = idc.NextHead(addr)
    return(False, None, None)


def run(struct_addr):
    struct_offsets = set([])
    out = {}
    xrefs_2_struct = get_to_xrefs(struct_addr)
    for addr in xrefs_2_struct:
        if idc.GetOperandValue(addr,1) != struct_addr:
            continue
        if idc.GetOpType(addr,0) != o_reg:
            continue
        reg = idc.GetOpnd(addr,0)
        next_addr = idc.NextHead(addr)
        status, offset, address =  forward_trace(next_addr, reg)
        if status:
            struct_offsets.add(offset)
    temp = list(sorted(struct_offsets))
    add_struct(temp)
    

# cursor at address of struct 
run(here())