IRDst = loc_key_2
@32[EBP_init + 0xFFFFFC2C] = 0x730077
@32[EBP_init + 0xFFFFFC20] = 0x770024
@32[EBP_init + 0xFFFFFC34] = 0x740062
@32[EBP_init + 0xFFFFFC28] = 0x6F0064
@32[EBP_init + 0xFFFFFC30] = 0x7E002E
@16[EBP_init + 0xFFFFFC38] = (EAX_init)[0:16]
@32[EBP_init + 0xFFFFFC24] = 0x6E0069

def run_symb_exec(symb, start, end):
    """ Execute symbolic execution and record memory writes

    :param symb         : SymbolicEngine
    :param start (int)  : start address
    :param end (int)    : end address
    return data (dict)  : Dictionnary with information needed (instruction offset, data_xrefs, source memory expression, spd value)
    """
    data = dict()
    while True:
        irblock = ircfg.get_block(start)
        if irblock is None:
            break
        for assignblk in irblock:
            if assignblk.instr:
                LOGGER.debug("0x%x : %s" % (assignblk.instr.offset, assignblk.instr))
                if assignblk.instr.args:
                    if assignblk.instr.args[0].is_mem() is True:  # write mem operation
                        data[assignblk.instr.args[0]] = {'to': assignblk.instr.offset, 'd_ref_type': 2, 'expr': assignblk.instr.args[0], 'spd': idc.get_spd(assignblk.instr.offset)}
                if assignblk.instr.offset == end:
                    break
            symb.eval_updt_assignblk(assignblk)
        start = symb.eval_expr(symb.ir_arch.IRDst)
    return data

def concretizes_stack_ptr(phrase_mem):
    """ replace ESP or EBP with a concrete value

    :param phrase_mem (dict)
    """
    stack = dict()
    for dst, src in phrase_mem.items():
        ptr_expr = dst.ptr
        spd = src['spd'] + 0x4
        real = expr_simp(ptr_expr.replace_expr({ExprId('ESP', 32): ExprInt(spd, 32), ExprId('EBP', 32): ExprInt(0x0, 32)}))
        if real.is_int():
            stack[real.arg] = src
    return stack

@128[ESP + 0x488] = {@8[0x423EE0] 0 8, @8[0x423EE0] ^ @128[0x423EE0][8:16] 8 16, @128[0x423EE0][16:128] 16 128}

{@8[0x423EE0] 0 8, @8[0x423EE0] ^ @128[0x423EE0][8:16] 8 16, @128[0x423EE0][16:128] 16 128} -> (((memory(0x423EE0, 0x1) & 0xff) << 0) | ((((memory(0x423EE0, 0x1) ^ ((memory(0x423EE0, 0x10) >> 8) & 0xff)) & 0xff) & 0xff) << 8) | ((((memory(0x423EE0, 0x10) >> 16) & 0xffffffffffffffffffffffffffff) & 0xffffffffffffffffffffffffffff) << 16))

def set_data_to_ida(start, data):
    """ Add comment with the strings, and redefine stack variables in IDA
    """
    size_local_var = idc.get_func_attr(start, idc.FUNCATTR_FRSIZE)
    frame_high_offset = 0x100000000
    frame_id = idc.get_func_attr(start, idc.FUNCATTR_FRAME)
    l_xref = dict()

    if size_local_var == 0:  # TODO : handle when stack size is undefined
        size_local_var = 0x1000  # set stack size to arbitrary value 0x1000

    # represent the stack as a list (initialization)
    stack = [0 for i in range(0, size_local_var)]

    # fill the stack
    for offset, value in data.items():
        if (offset >> 31) == 1:  # EBP based frame
            if offset <= frame_high_offset:
                stack_offset = twos_complement(offset)  # Ex: 0xfffffe00 -> 0x200
                frame_offset = stack_to_frame_offset(stack_offset, size_local_var)  # frame_offset = member_offset = positive offset in the stack frame
                l_xref[frame_offset] = value
                stack[frame_offset] = value['value']
        else:  # ESP based frame
            l_xref[offset] = value
            stack_offset = frame_to_stack_offset(offset, size_local_var)
            frame_offset = offset
            stack[offset] = value['value']
        LOGGER.debug("offset 0x%x, stack_offset 0x%x, frame_offset 0x%x, value %s" % (offset, stack_offset, frame_offset, chr(value['value'])))

    index = 0
    while index < len(stack):
        if stack[index] != 0:  # skip null bytes
            buff = get_bytes_until_delimiter(stack[index::], [0x0, 0x0])  # TODO: adjust auto ascii/utf16 ?

            if buff:
                LOGGER.debug("raw output = %s" % repr(buff))
                ascii_buff = zeroes_out(buff)  # transform to ascii, deletes null bytes

                if buff[1] == 0x0:  # check if strings is utf16 TODO: check if really usefull
                    STRTYPE = idc.STRTYPE_C
                else:
                    STRTYPE = idc.STRTYPE_C

                var_offset = index
                var_name = idc.get_member_name(frame_id, var_offset)
                s_string = "".join(map(chr, ascii_buff))
                LOGGER.info("[+] Index 0x%x, var_offset 0x%x, frame_id = 0x%x, var_name %s, strings = %s, clean_string = %s, len_strings = 0x%x, size_local_var 0x%x, xref to 0x%x" % (index, var_offset, frame_id, var_name, s_string, replace_forbidden_char(s_string), len(buff), size_local_var, l_xref[index]['to']))

                # delete struct member, to create array of char
                LOGGER.info("[+] Delete structure member")
                for a_var_offset in range(var_offset, var_offset + len(buff)):
                    # if idc.get_member_name(frame_id, a_var_offset):
                    ret = idc.del_struc_member(frame_id, a_var_offset)
                    LOGGER.debug("ret del_struc_member = %d" % ret)

                # re add structure member with good size and proper name
                LOGGER.info("[+] Add structure member")
                ret = idc.add_struc_member(frame_id, replace_forbidden_char(s_string)[:0x10], var_offset, idc.FF_STRLIT, STRTYPE, len(buff))
                LOGGER.debug("ret add_struc_member = %d" % ret)

                # Comment instruction using the xref key
                LOGGER.info("[+] Comment instruction")
                if l_xref[index]['d_ref_type'] == 2:
                    if not idc.get_cmt(l_xref[index]['to'], 0):
                        idc.set_cmt(l_xref[index]['to'], s_string, 0)
                index = index + len(buff)
            else:
                index = index + 1
        else:
            index = index + 1
    LOGGER.info("[+] end")
    return stack

rule malware_first_unpacking_routine {
    strings :
        $h1 = { 81 [1-5] A9 0F 00 00 75 ?? C7 ?? ?? ?? ?? ?? 40 2E EB ED }
    condition :
        $h1
}

import idaapi
import idc
import logging
from math import log


def unpack_arbitrary(data, word_size=None):
    """ (modified pwntools) unpack arbitrary-sized integer
    :param data     : String to convert
    :word_size      : Word size of the converted integer or the string “all” (in bits).
    :return        : The unpacked number
    """
    number = 0
    data = reversed(data)
    data = bytearray(data)

    for c in data:
        number = (number << 8) + c

    number = number & ((1 << word_size) - 1)

    signbit = number & (1 << (word_size-1))
    return int(number - 2*signbit)


def memory(ea, size):
    """ Read memory from IDA
    :param ea (int)     : Start address
    :param size (int)   : size of buffer in normal 8-bit bytes
    : return int/long
    """
    return unpack_arbitrary(idc.get_bytes(ea, size), size * 8)


def twos_complement(esp_offset):
    return 0x100000000 - esp_offset


def stack_to_frame_offset(stack_offset, size_local_var):
    """
    stack_offset <-> frame_offset
    """
    return size_local_var - abs(stack_offset)


def frame_to_stack_offset(frame_offset, size_local_var):
    """
    stack_offset <-> frame_offset
    """
    return size_local_var - abs(frame_offset)


def get_bytes_until_delimiter(bytes_list, delimiter):
    """Get bytes until it reach the delimiter

    :param bytes_list (list)    : List of bytes
    :param delimiter (list)     : Delimiter to reach (e.g. [0x0, 0x0])
    :return a list of bytes
    """

    len_delimiter = len(delimiter)
    iterables = list()
    for x in range(len_delimiter):
        iterables.append(bytes_list[x:])
    abc = zip(*iterables)
    if tuple(delimiter) in abc:
        index = abc.index(tuple(delimiter))
        return (bytes_list[0:index]) + delimiter
    else:
        return None


def zeroes_out(bytes_list):
    """ delete null bytes

    :param bytes_list (list)    : list of bytes
    :return                     : list of bytes without zeros
    """
    return [x for x in bytes_list if x != 0]


def replace_forbidden_char(my_strings):
    """ replace IDA forbidden char
    """

    clean_string = list()
    for x in my_strings:
        if x.isalnum():
            clean_string.append(x)
        elif x.isspace():
            clean_string.append('_')
        else:
            clean_string.append('?')
    return "k_" + "".join(clean_string)


def bytes_needed(n):
    """ get the number of bytes that compose the number
    https://stackoverflow.com/questions/14329794/get-size-of-integer-in-python

    :pararm n (int) : number
    :return (int)   : number of bytes
    """
    if n == 0:
        return 1
    return int(log(n, 256)) + 1


def extract_byte_x(num, position):
    """ Get bytes from position

    :param num (int)        : original number
    :param position (int)   : desired byte position
    : return (int)          : byte x
    """
    offset = 0x8 * (position - 1)
    and_mask = (0xFF << offset)
    return (num & and_mask) >> offset


def symbolic_exec(start, end):
    """ Symbolic execution engine (modified of original https://github.com/cea-sec/miasm/blob/master/example/ida/symbol_exec.py)

    Takes start and end address from IDA selection, do symbolic execution on instructions
    Replace stacks registers (ESP, EBP) variable by concrete value 0x0 or current spd instructions
    """
    from miasm.ir.symbexec import SymbolicExecutionEngine
    from miasm.expression.expression import ExprId, ExprInt
    from miasm.expression.simplifications import expr_simp
    from miasm.core.bin_stream_ida import bin_stream_ida
    from miasm.ir.translators import Translator
    from miasm.analysis.machine import Machine

    def run_symb_exec(symb, start, end):
        """ Execute symbolic execution and record memory writes

        :param symb         : SymbolicEngine
        :param start (int)  : start address
        :param end (int)    : end address
        return data (dict)  : Dictionnary with information needed (instruction offset, data_xrefs, source memory expression, spd value)
        """
        data = dict()
        while True:
            irblock = ircfg.get_block(start)
            if irblock is None:
                break
            for assignblk in irblock:
                if assignblk.instr:
                    LOGGER.debug("0x%x : %s" % (assignblk.instr.offset, assignblk.instr))
                    if assignblk.instr.args:
                        if assignblk.instr.args[0].is_mem() is True:  # write mem operation
                            data[assignblk.instr.args[0]] = {'to': assignblk.instr.offset, 'd_ref_type': 2, 'expr': assignblk.instr.args[0], 'spd': idc.get_spd(assignblk.instr.offset)}
                    if assignblk.instr.offset == end:
                        break
                symb.eval_updt_assignblk(assignblk)
            start = symb.eval_expr(symb.ir_arch.IRDst)
        return data

    def concretizes_stack_ptr(phrase_mem):
        """ replace ESP or EBP with a concrete value

        :param phrase_mem (dict)
        """
        stack = dict()
        for dst, src in phrase_mem.items():
            ptr_expr = dst.ptr
            spd = src['spd'] + 0x4
            real = expr_simp(ptr_expr.replace_expr({ExprId('ESP', 32): ExprInt(spd, 32), ExprId('EBP', 32): ExprInt(0x0, 32)}))
            if real.is_int():
                stack[real.arg] = src
        return stack

    def eval_src_expr(symb, data):
        """ evaluate/resolve the source expression trying to get a concrete value.
        Translate to python anbd eval if necessary

        Evaluation:
        @128[ESP + 0x488] = {@8[0x423EE0] 0 8, @8[0x423EE0] ^ @128[0x423EE0][8:16] 8 16, @128[0x423EE0][16:128] 16 128}
        Translation:
        {@8[0x423EE0] 0 8, @8[0x423EE0] ^ @128[0x423EE0][8:16] 8 16, @128[0x423EE0][16:128] 16 128} = (((memory(0x423EE0, 0x1) & 0xff) << 0) | ((((memory(0x423EE0, 0x1) ^ ((memory(0x423EE0, 0x10) >> 8) & 0xff)) & 0xff) & 0xff) << 8) | ((((memory(0x423EE0, 0x10) >> 16) & 0xffffffffffffffffffffffffffff) & 0xffffffffffffffffffffffffffff) << 16))
        """
        tmp = dict()
        for dst, src in data.items():
            tmp_val = symb.eval_expr(src['expr'])
            if tmp_val.is_int():  # Case where the source expression is a int
                tmp_val = int(tmp_val.arg)
            else:  # The source expression may by a data reference (read from mem)
                # Translate expression to python
                py_expr = Translator.to_language('Python').from_expr(tmp_val)
                try:
                    tmp_val = eval(py_expr)  # TODO better solution (need to add memory function) ?
                except Exception as e:
                    LOGGER.info("[-] Python Translator Exception")
                    LOGGER.info(py_expr)
                    raise e
            if isinstance(tmp_val, int) or isinstance(tmp_val, long):
                # Set one byte per address
                for i in range(0, bytes_needed(tmp_val)):
                    tmp[dst + i] = {'to': src['to'], 'd_ref_type': 2, 'value': extract_byte_x(tmp_val, i + 1), 'spd': src['spd']}
        return tmp

    bs = bin_stream_ida()
    machine = Machine("x86_32")
    mdis = machine.dis_engine(bs, dont_dis_nulstart_bloc=True)

    ir_arch = machine.ira(mdis.loc_db)

    # Disassemble the targeted function until end
    mdis.dont_dis = [end]
    asmcfg = mdis.dis_multiblock(start)

    # Generate IR
    ircfg = ir_arch.new_ircfg_from_asmcfg(asmcfg)

    # Replace ExprID
    regs = machine.mn.regs.regs_init
    # regs[ExprId('ESP', 32)] = ExprId('ESP', 32)
    # regs[ExprId('ESP', 32)] = ExprId('EBP', 32)
    # regs[ExprId('EAX', 32)] = ExprId('EAX', 32)
    # regs[ExprId('EBX', 32)] = ExprId('EBX', 32)
    # regs[ExprId('ECX', 32)] = ExprId('ECX', 32)
    # regs[ExprId('EDX', 32)] = ExprId('EDX', 32)
    # regs[ExprId('XMM0_init', 32)] = ExprId('XMM0', 32)

    LOGGER.info("[+] Get symbolic engine")
    symb = SymbolicExecutionEngine(ir_arch, regs)
    LOGGER.info("[+] Run symbolic execution")
    data = run_symb_exec(symb, start, end)
    LOGGER.info("[+] Concretize stack pointer")
    data = concretizes_stack_ptr(data)
    LOGGER.info("[+] Resolves source expression")
    data = eval_src_expr(symb, data)
    return data


def set_data_to_ida(start, data):
    """ Add comment with the strings, and redefine stack variables in IDA
    """
    size_local_var = idc.get_func_attr(start, idc.FUNCATTR_FRSIZE)
    frame_high_offset = 0x100000000
    frame_id = idc.get_func_attr(start, idc.FUNCATTR_FRAME)
    l_xref = dict()

    if size_local_var == 0:  # TODO : handle when stack size is undefined
        size_local_var = 0x1000  # set stack size to arbitrary value 0x1000

    # represent the stack as a list (initialization)
    stack = [0 for i in range(0, size_local_var)]

    # fill the stack
    for offset, value in data.items():
        if (offset >> 31) == 1:  # EBP based frame
            if offset <= frame_high_offset:
                stack_offset = twos_complement(offset)  # Ex: 0xfffffe00 -> 0x200
                frame_offset = stack_to_frame_offset(stack_offset, size_local_var)  # frame_offset = member_offset = positive offset in the stack frame
                l_xref[frame_offset] = value
                stack[frame_offset] = value['value']
        else:  # ESP based frame
            l_xref[offset] = value
            stack_offset = frame_to_stack_offset(offset, size_local_var)
            frame_offset = offset
            stack[offset] = value['value']
        LOGGER.debug("offset 0x%x, stack_offset 0x%x, frame_offset 0x%x, value %s" % (offset, stack_offset, frame_offset, chr(value['value'])))

    index = 0
    while index < len(stack):
        if stack[index] != 0:  # skip null bytes
            buff = get_bytes_until_delimiter(stack[index::], [0x0, 0x0])  # TODO: adjust auto ascii/utf16 ?

            if buff:
                LOGGER.debug("raw output = %s" % repr(buff))
                ascii_buff = zeroes_out(buff)  # transform to ascii, deletes null bytes

                if buff[1] == 0x0:  # check if strings is utf16 TODO: check if really usefull
                    STRTYPE = idc.STRTYPE_C_16
                else:
                    STRTYPE = idc.STRTYPE_C

                var_offset = index
                var_name = idc.get_member_name(frame_id, var_offset)
                s_string = "".join(map(chr, ascii_buff))
                LOGGER.info("[+] Index 0x%x, var_offset 0x%x, frame_id = 0x%x, var_name %s, strings = %s, clean_string = %s, len_strings = 0x%x, size_local_var 0x%x, xref to 0x%x" % (index, var_offset, frame_id, var_name, s_string, replace_forbidden_char(s_string), len(buff), size_local_var, l_xref[index]['to']))

                # delete struct member, to create array of char
                LOGGER.info("[+] Delete structure member")
                for a_var_offset in range(var_offset, var_offset + len(buff)):
                    # if idc.get_member_name(frame_id, a_var_offset):
                    ret = idc.del_struc_member(frame_id, a_var_offset)
                    LOGGER.debug("ret del_struc_member = %d" % ret)

                # re add structure member with good size and proper name
                LOGGER.info("[+] Add structure member")
                ret = idc.add_struc_member(frame_id, replace_forbidden_char(s_string)[:0x10], var_offset, idc.FF_STRLIT, STRTYPE, len(buff))
                LOGGER.debug("ret add_struc_member = %d" % ret)

                # Comment instruction using the xref key
                LOGGER.info("[+] Comment instruction")
                if l_xref[index]['d_ref_type'] == 2:
                    if not idc.get_cmt(l_xref[index]['to'], 0):
                        idc.set_cmt(l_xref[index]['to'], s_string, 0)
                index = index + len(buff)
            else:
                index = index + 1
        else:
            index = index + 1
    LOGGER.info("[+] end")
    return stack


def main():
    start, end = idc.read_selection_start(), idc.read_selection_end()

    if start != 0xFFFFFFFF and end != 0xFFFFFFFF:
        data = symbolic_exec(start, end)
        LOGGER.info("[+] Set and comment data to IDA")
        stack = set_data_to_ida(start, data)
    else:
        print("Error: Select instructions")
        stack = -0x1
    return stack


if __name__ == '__main__':
    idaapi.CompileLine('static key_F3() { RunPythonStatement("main()"); }')
    idc.add_idc_hotkey("F3", "key_F3")

    LOGGER = logging.getLogger(__name__)
    if not LOGGER.handlers:  # Avoid duplicate handler
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(logging.Formatter("[%(levelname)s]  %(message)s"))
        LOGGER.addHandler(console_handler)
        LOGGER.setLevel(logging.DEBUG)

    print("=" * 50)
    print("""Available commands:
    main() - F3: Symbolic execution of current selection
    """)