import collections import struct import sys class BitArray(object): def __init__(self, size, bits): self.size = size if len(bits) > size: raise ValueError("size > len(bits)") bits_list = [] for bit in bits: x = int(bit) if x not in [0, 1]: raise ValueError("Not expected bits value {0}".format(x)) bits_list.append(x) self.array = bits_list if size > len(self.array): self.array = ([0] * (size - len(self.array))) + self.array def dump(self): res = [] for i in range(self.size // 8): c = 0 for x in (self.array[i * 8: (i + 1) * 8]): c = (c << 1) + x res.append(c) return bytearray((res)) def __getitem__(self, slice): return self.array[slice] def __setitem__(self, slice, value): self.array[slice] = value return True def __repr__(self): return repr(self.array) def __add__(self, other): if not isinstance(other, BitArray): return NotImplemented return BitArray(self.size + other.size, self.array + other.array) def to_int(self): return int("".join([str(i) for i in self.array]), 2) @classmethod def from_string(cls): l = [] for c in bytearray(reversed(str_base)): for i in range(8): l.append(c & 1) c = c >> 1 self.array = l @classmethod def from_int(cls, size, x): if x < 0: x = x & ((2 ** size) - 1) return cls(size, bin(x)[2:]) # Rules: bytes only !!!! mem_access = collections.namedtuple('mem_access', ['base', 'index', 'squale', 'disp']) x86_regs = ['EAX', 'ECX', 'EDX', 'EBX', 'ESP', 'EBP', 'ESI', 'EDI'] def create_displacement(base=None, index=None, squale=None, disp=0): return mem_access(base, index, squale, disp) class X86RegisterSelector(object): size = 3 # bits reg_order = ['EAX', 'ECX', 'EDX', 'EBX', 'ESP', 'EBP', 'ESI', 'EDI'] reg_opcode = {v : BitArray.from_int(size=3, x=i) for i, v in enumerate(reg_order)} def accept_arg(self, previous, args): x = args[0] try: return (1, self.reg_opcode[x]) except KeyError: return (None, None) @classmethod def get_reg_bits(cls, name): return cls.reg_opcode[name] class RawBits(BitArray): def accept_arg(self, previous, args): return (0, self) class Imm32(object): def accept_arg(self, previous, args): x = int(args[0]) return (1, BitArray.from_int(32, X86.to_little_endian(x))) class ModRM(object): size = 8 def __init__(self, *sub_modrm): self.sub = sub_modrm def accept_arg(self, previous, args): if len(args) < 2: raise ValueError("Missing arg for modrm") arg1 = args[0] arg2 = args[1] for sub in self.sub: #import pdb;pdb.set_trace() if sub.match(arg1, arg2): d = sub(arg1, arg2, 0) previous[0][-2] = d.direction return (2, d.mod + d.reg + d.rm + d.after) elif sub.match(arg2, arg1): d = sub(arg2, arg1, 1) previous[0][-2] = d.direction return (2, d.mod + d.reg + d.rm + d.after) return (None, None) class X86(object): @staticmethod def is_reg(name): return name in x86_regs @staticmethod def is_mem_acces(data): return isinstance(data, mem_access) @staticmethod def mem_access_has_only(mem_access, names): if not X86.is_mem_acces(mem_access): raise ValueError("mem_access_has_only") for f in mem_access._fields: if getattr(mem_access, f) and f not in names: return False return True @staticmethod def to_little_endian(i): i = i & 0xffffffff return struct.unpack("I", i))[0] class ModRM_REG__REG(object): @classmethod def match(cls, arg1, arg2): return X86.is_reg(arg1) and X86.is_reg(arg2) def __init__(self, arg1, arg2, reversed): self.mod = BitArray(2, "11") self.reg = X86RegisterSelector.get_reg_bits(arg2) self.rm = X86RegisterSelector.get_reg_bits(arg1) self.after = BitArray(0, "") self.direction = 0 class ModRM_REG__DEREF_REG(object): @classmethod def match(cls, arg1, arg2): return X86.is_reg(arg1) and arg1 not in ["ESP", "EBP"] and X86.is_mem_acces(arg2) and X86.mem_access_has_only(arg2, ["base"]) def __init__(self, arg1, arg2, reversed): self.mod = BitArray(2, "00") self.reg = X86RegisterSelector.get_reg_bits(arg1) self.rm = X86RegisterSelector.get_reg_bits(arg2.base) self.after = BitArray(0, "") self.direction = not reversed class ModRM_REG__DEREF_REG_IMM(object): @classmethod def match(cls, arg1, arg2): return X86.is_reg(arg1) and X86.is_mem_acces(arg2) and X86.mem_access_has_only(arg2, ["base", "disp"]) def __init__(self, arg1, arg2, reversed): self.mod = BitArray(2, "10") self.reg = X86RegisterSelector.get_reg_bits(arg1) self.rm = X86RegisterSelector.get_reg_bits(arg2.base) self.after = BitArray.from_int(32, X86.to_little_endian(arg2.disp)) self.direction = not reversed class ModRM_REG_IMM(object): @classmethod def match(cls, arg1, arg2): return arg1 in x86_regs and arg2 in x86_regs def __init__(self, arg1, arg2): self.mod = BitArray(2, "11") self.reg = X86RegisterSelector.get_reg_bits(arg2) self.rm = X86RegisterSelector.get_reg_bits(arg1) self.direction = 0 class Instruction(object): encoding = [] def __init__(self, *initial_args): for type_encoding in self.encoding: args = list(initial_args) res = [] for element in type_encoding: arg_consum, value = element.accept_arg(res, args) if arg_consum is None: break res.append(value) del args[:arg_consum] else: # if no break if args: # if still args: fail continue self.value = sum(res, BitArray(0, "")) return raise ValueError("Cannot encode :(") class Push(Instruction): encoding = [(RawBits.from_int(5, 0x50 >> 3), X86RegisterSelector()), (RawBits.from_int(8, 0x68), Imm32())] class Pop(Instruction): encoding = [(RawBits.from_int(5, 0x58 >> 3), X86RegisterSelector())] class Mov(Instruction): encoding = [(RawBits.from_int(8, 0x89), ModRM(ModRM_REG__REG, ModRM_REG__DEREF_REG, ModRM_REG__DEREF_REG_IMM)), (RawBits.from_int(5, 0xb8 >> 3), X86RegisterSelector(), Imm32())] class Call(Instruction): encoding = [(RawBits.from_int(13, 0xffd0 >> 3), X86RegisterSelector())] class Ret(Instruction): encoding = [(RawBits.from_int(8, 0xc3),)] class MultipleInstr(object): def __init__(self): self.instrs = [] def __iadd__(self, value): if type(value) == MultipleInstr: self.instrs.extend(value.instrs) return self self.instrs.append(value) return self def get_code(self): if sys.version_info.major == 3: return b"".join([x.value.dump() for x in self.instrs]) return "".join([str(x.value.dump()) for x in self.instrs])