idaware.py

#!/usr/bin/python3

from functools import wraps
import inspect

import json

import idaapi
import idautils
import ida_funcs
import idc

from capstone.x86 import *
from capstone import *

DEBUG = True

def symbolAddr(name): return [tup for tup in idautils.Names() if tup[1] == name][0][0]

import_idaHook = """#autogenerated hook code
def idaHook(addr, enabled=True, halting=False):
    def init(func):
        @wraps(func)
        def wrapper():
            if not enabled: return False
            if halting == None:
                return func()
            func()
            return halting
        return wrapper
    return init


"""

def idaHook(addr, enabled=True, halting=False):
    
    def init(func):

        idaapi.add_bpt(addr, enabled, idaapi.BPT_SOFT)

        @wraps(func)
        def wrapper():
            # let the hook determine
            if halting == None: return func()
            func()
            return halting

        bpt = idaapi.bpt_t()
        idaapi.get_bpt(addr, bpt)
        bpt.elang = 'Python'

        cond = import_idaHook + inspect.getsource(wrapper) + "\n" + wrapper.__name__ + "()"
        bpt.condition = cond
        idaapi.update_bpt(bpt)

        return wrapper
    return init


def fmtHex(hexArr, fmt='{hexByte} '):

    retStr = ''
    for byte in hexArr:
        byteStr = hex(byte)[2:].rjust(2, "0")
        retStr += fmt.format(hexByte = byteStr)

    return retStr


jmps = [X86_INS_JMP,
    X86_INS_JAE,
    X86_INS_JA,
    X86_INS_JBE,
    X86_INS_JB,
    X86_INS_JECXZ,
    X86_INS_JE,
    X86_INS_JGE,
    X86_INS_JG,
    X86_INS_JLE,
    X86_INS_JL,
    X86_INS_JNE,
    X86_INS_JNO,
    X86_INS_JNP,
    X86_INS_JNS,
    X86_INS_JO,
    X86_INS_JP,
    X86_INS_JRCXZ,
    X86_INS_JS,
    X86_INS_CALL]

def makeSig(name, siglen):

    code = idaapi.get_bytes(symbolAddr(name), siglen)

    md = Cs(CS_ARCH_X86, CS_MODE_64)
    md.detail = True

    mask = ''
    trackedLen = 0

    for i in md.disasm(code, 0):

        # if i.id == X86_INS_CMP: return i

        trackedMask = ''

        if DEBUG: 
            trackedBytes = b''
            # print("insn:", i.mnemonic + " " + i.op_str)

        if DEBUG: prefixBytes = b''
        for prefixByte in i.prefix:
            if prefixByte:
                # sig += prefixByte.to_bytes(1, sys.byteorder)
                trackedMask += 'x'
                mask += 'x'
                if DEBUG:
                    prefixBytes += prefixByte.to_bytes(1, sys.byteorder)

        if prefixBytes and DEBUG:
            # print("PREFIX:", fmtHex(prefixBytes))
            trackedBytes += prefixBytes

        if i.rex:
            # sig += i.rex.to_bytes(1, sys.byteorder) # byteorder is important here
            trackedMask += 'x'
            mask += 'x'
            if DEBUG:
                # print("REX:", hex(i.rex))
                trackedBytes += i.rex.to_bytes(1, sys.byteorder)

        if DEBUG: opcodeBytes = b''
        for opcodeByte in i.opcode:
            if opcodeByte:
                # sig += opcodeByte.to_bytes(1, sys.byteorder)
                trackedMask += 'x'
                mask += 'x'
                if DEBUG:
                    opcodeBytes += opcodeByte.to_bytes(1, sys.byteorder)

        if opcodeBytes and DEBUG:
            # print("opcode:", fmtHex(opcodeBytes))
            trackedBytes += opcodeBytes

        if i.modrm:
            # sig += i.modrm.to_bytes(1, sys.byteorder)
            trackedMask += 'x'
            mask += 'x'
            if DEBUG:
                # print("MODRM:", hex(i.modrm))
                trackedBytes += i.modrm.to_bytes(1, sys.byteorder)

        if i.sib:
            # sig += i.sib.to_bytes(1, sys.byteorder)
            trackedMask += 'x'
            mask += 'x'
            if DEBUG:
                # print("SIB:", hex(i.sib))
                trackedBytes += i.sib.to_bytes(1, sys.byteorder)
            
        if i.disp_size:
            for x in range(i.disp_size):
                # sig += b'\x00'
                mask += '?'
                trackedMask += '?'
                if DEBUG:
                    trackedBytes += b'\x00'

        jmp = False
        if i.id in jmps:
            for x in range(i.imm_size):
                # sig += b
                mask += '?'
                trackedMask += '?'
                if DEBUG:
                    trackedBytes += b'\x00'

        else:
            trackedBytes += i.imm_offset.to_bytes(i.imm_size, sys.byteorder)
            mask += 'x' * i.imm_size
            trackedMask += 'x' * i.imm_size
            # sig += i.imm_offset.to_bytes(i.imm_size, sys.byteorder)


        # these two should be the same
        if DEBUG:
            print("insn mask:", trackedMask)
            print("new bytes:", fmtHex(trackedBytes))
            print("old bytes:", fmtHex(bytes(i.bytes)))
    
    code = code[0:len(mask)]
    return {name: {"sig": fmtHex(code)[0:-1], "mask": mask}}

@unwrapName := lambda f: (lambda sigmask: f(sigmask, list(sigmask)[0]))
def printSig(sigmask, name):

    outStr = ''
    sig  = sigmask[name]["sig"]
    mask = sigmask[name]["mask"] 

    sigBytes = sig.split(' ')
    
    for x in range(len(sigBytes)):
        if(mask[x] == 'x'):
            outStr += sigBytes[x] + ' '
        else:
            outStr += '?? '
 
    print(outStr)

print("DONE")  

def getFuncDetails(name):

    addr = symbolAddr(name)
    funcSize = ida_funcs.get_func(addr).size()

    return json.dumps({
        "name":name,
        "addr":hex(addr),
        "bytes": fmtHex(idaapi.get_bytes(addr, funcSize), fmt='\\x{hexbyte}')[0:-1]
        }, indent=2)