idaisn

import idautils
from idaapi import *
from idc import *
import enum
from time import sleep
import ctypes

kOperandNone = 0
kOperandReg = 1
kOperandMemory = 2
kOperandImmediate = 3
kOperandBranch = 4

kSegNone = 0
kSegGS = 1
kSegFS = 2
kSegDS = 3
kSegSS = 4

class Operand:
    def __init__(self, otype, base, index, scale, disp, imm, seg):
        self.type = otype
        self.seg = seg
        self.base = base

    def seg(self):
        return self.seg
    def base_reg(self):
        return self.base
    def type(self):
        return self.type

class Instruction:
    def __init__(self, addr, length):
        self.mnem = idc.print_insn_mnem(addr)
        self.address = addr
        self.operands = []
        self.line = idc.generate_disasm_line(addr, 1)
        self.opcodes = ida_bytes.get_bytes(addr, length)
    def type(self):
        return self.mnem
    def add_operand(self, operand):
        self.operands.append(operand)
    def text(self):
        return self.line
    def operands(self):
        return self.operands
    def address(self):
        return self.address
    def opcodes(self):
        return self.opcodes

class BasicBlock:
    def __init__(self, ea):
        self.start = ea
        self.instructions = []
    def add_instruction(self, i):
        self.instructions.append(i)
    def address(self):
        return self.start
    def instructions(self):
        return self.instructions

class Function:
    def __init__(self, addr):
        self.addr = addr
        self.blocks = []
    def add_block(self, bb):
        self.blocks.append(bb)
    def blocks(self):
        return self.blocks

def save_instruction(ea):
    #print("Save Instruction")
    insn = idaapi.insn_t()
    insnLen = idaapi.decode_insn(insn, ea)
    i = Instruction(ea, insnLen)
    #print(i.type())
    DATA_SIZE = {dt_byte: 1, dt_word: 2, dt_dword: 4, dt_tbyte: 4, dt_float: 4, dt_qword: 8, dt_double: 8, dt_byte16: 16}
    SEG_NAMES = {0: kSegNone, 29: kSegNone, 31: kSegSS, 32: kSegDS, 34: kSegGS}
    for oidx in range(8):
        operand = insn.ops[oidx]
        segrg  = (operand.specval & 0xffff0000) >> 16
        segsel = (operand.specval & 0x0000ffff) >> 0
        #print("---")
        #print(operand.type)
        #print(f"segreg: {segrg} segsel: {segsel}")


        if operand.type == ida_ua.o_reg:
            reg = operand.reg
            #print(f"reg: {get_reg_name(reg,DATA_SIZE[operand.dtype])}")
            tmp = Operand(kOperandReg, get_reg_name(reg,DATA_SIZE[operand.dtype]), None, None, None, None, kSegNone)
            i.add_operand(tmp)
            #print("Saved - kOperandReg")
        elif operand.type == ida_ua.o_displ or operand.type == ida_ua.o_phrase or operand.type == ida_ua.o_mem:
            sib = operand.specflag2
            hasSIB = operand.specflag1
            #print("HAS SIB: " + str(hasSIB))
            base = sib  & 7
            #print(f"{hex(ea)} - " + idc.generate_disasm_line(ea, 1))
            base_size = DATA_SIZE[operand.dtype]
            index = (sib >> 3) & 7
            index_size = DATA_SIZE[operand.dtype]
            scale = (sib >> 6) & 3
            disp = operand.addr
            reg = operand.reg
            #print(f"{hex(ea)} - " + idc.generate_disasm_line(ea, 1))
            seg = SEG_NAMES[segrg]

            #print(f"reg: {get_reg_name(reg,base_size)} base: {get_reg_name(base,base_size)} idx: {get_reg_name(index,index_size)} scale: {scale*2} disp: {hex(disp)}")

            if not hasSIB:
                reg = get_reg_name(reg,base_size)
                if operand.type == ida_ua.o_mem:
                    reg = None
                tmp = Operand(kOperandMemory, reg, None, None, disp, None, seg)
                i.add_operand(tmp)
            else:
                if operand.type == ida_ua.o_mem:
                    tmp = Operand(kOperandMemory, get_reg_name(index,index_size), None, scale*2, disp, None, seg)
                    i.add_operand(tmp)
                else:
                    tmp = Operand(kOperandMemory, get_reg_name(base,base_size), get_reg_name(index,index_size), scale*2, disp, None, seg)
                    i.add_operand(tmp)

            #print("Saved - kOperandMemory")
        elif operand.type == ida_ua.o_imm:
            value = operand.value
            #print(f"value: {hex(value)}")
            tmp = Operand(kOperandImmediate, None, None, None, None, value, kSegNone)
            i.add_operand(tmp)
            #print("Saved - kOperandImmediate")
        elif operand.type == ida_ua.o_near:
            value = operand.addr
            #print(f"value: {hex(value)}")
            tmp = Operand(kOperandBranch, None, None, None, None, value, kSegNone)
            i.add_operand(tmp)
            #print("Saved - kOperandBranch")
        else:
            tmp = Operand(kOperandNone, None, None, None, None, None, None)
            i.add_operand(tmp)
    #print("Saved")
    return i