IDA script for deobfuscation of Nymaim malware

deobfuscation.py

import idaapi

idaapi.CompileLine('static deobfuscate() { RunPythonStatement("deobfuscate()"); }')
AddHotkey("Alt-N", "deobfuscate")

repl_pairs = {
    "e8 3a 00 00 00" : "b8 01 00 00 00", # mov eax, 1
}

def chunk(string, length):
    return [string[i:i+length:] for i in range(0, len(string), length)]

def pad(str):
    if len(str) == 1:
        str = "0" + str
    return str

def get_byte(ea):
    return pad(hex(Byte(ea))[2::])

def get_n_bytes(ea, n):
    bytes = []
    for i in range(n):
        bytes.append(get_byte(ea + i))
    return " ".join(bytes)

def get_ins_bytes(ea):
    bytes = []
    for i in range(ItemSize(ea)):
        bytes.append(get_byte(ea + i))
    return " ".join(bytes)

def get_n_ins_bytes(ea, n):
    bytes = []
    addr = ea
    for i in range(n):
        bytes.append(get_ins_bytes(addr))
        addr = NextHead(addr)
    return " ".join(bytes)

def nop(start, end):
    addr = start
    while addr < end:
        PatchByte(addr, 0x90) # NOP
        addr += 1
    HideArea(start, end, '', '', '', 0)

def warn_overflow():
    print "Replacement is longer than original instruction. Proceed? (y/n)"
    resp = raw_input().lower()
    return resp == 'y'

def replace(start, replacement, len_to_replace):
    repl = replacement.split()
    end = start + len_to_replace
    for i in range(len_to_replace):
        addr = start + i
        if i >= len(repl):
            nop(addr, end)
            return
        PatchByte(addr, int(repl[i], 16))
    Refresh()
    print "Code replaced."

def is_mnem(start, mnem):
    return GetMnem(start) == mnem

def is_call_x(start, x):
    return is_mnem(start, 'call') and GetOpnd(start, 0) == x

def get_push_register(start):
    op_reg = {'4F':'eax','50':'ecx','51':'edx','52':'ebx','54':'ebp','55':'esi','56':'edi'}
    if is_mnem(start, 'push'):
        operand = GetOpnd(start, 0)[:-1:]
        nxt = NextHead(start)
        if operand in op_reg and is_call_x(nxt, 'set_top_stack_element'):
            return op_reg[operand]
    return None

def add(a, b):
    return b + a

def sub(a, b):
    return b - a
    
def xor(a, b):
    return a ^ b

def do_op(first_opr, sec_opr, op):
    op_map = {'add' : add, 'sub' : sub, 'xor' : xor}
    return op_map[op](first_opr, sec_opr)

def get_fn_addr(start):
    addr_2 = NextHead(start)
    addr_3 = NextHead(addr_2)
    addr_4 = NextHead(addr_3)
    addr_ret = NextHead(addr_4)
    if is_mnem(start, 'push') and is_mnem(addr_2, 'push') and is_mnem(addr_3, 'push') and is_mnem(addr_4, 'call'):
        fn_name = GetOpnd(addr_4, 0)
        if 'special_jump' in fn_name:
            first_opr = int(GetOpnd(addr_3, 0)[:-1:], 16)
            sec_opr = int(GetOpnd(addr_2, 0)[:-1:], 16)
            op = fn_name.split("_")[2][:3:]
            return (addr_ret + do_op(first_opr, sec_opr, op)) & 0xffffffff
    return None
    
def gen_call_bytes(start, fn_addr):
    ins = 'call ' + str(fn_addr)
    machine_bytes = Assemble(start, ins)[1].encode("hex")
    return " ".join(chunk(machine_bytes, 2))

def get_push_code(reg):
    reg_to_code = {
        "eax" : "50",
        "ecx" : "51",
        "edx" : "52",
        "ebx" : "53",
        "ebp" : "55",
        "esi" : "56",
        "edi" : "57",
    }
    return reg_to_code[reg]

def get_replacement(start):
    # register push
    reg = get_push_register(start)
    if reg:
        return get_n_ins_bytes(start, 2), get_push_code(reg)

    # obfuscated fn call
    fn_addr = get_fn_addr(start)
    if fn_addr:
        return get_n_ins_bytes(start, 4), gen_call_bytes(start, fn_addr)
    
    # arbitrary replacement
    candidates = map(lambda keys: keys.split(), repl_pairs.keys())
    i = 0
    while len(candidates) > 0:
        if len(candidates) == 1:
            joined_candidate = " ".join(candidates[0])
            if get_n_bytes(start, len(candidates[0])) == joined_candidate:
                return joined_candidate, repl_pairs[joined_candidate]
            return None, None
        candidates = filter(lambda candidate: candidate[i] == get_byte(start + i), candidates)
        i += 1

    # no match found
    return None, None

def deobfuscate():
    start = ScreenEA()
    end = NextHead(start)
    candidate, replacement = get_replacement(start)
    if replacement:
        if len(replacement) > len(candidate) and not warn_overflow():            
            print "Instruction not replaced"
            return
        replace(start, replacement, len(candidate.split(" ")))
        return
    print "No match found."