diff --git a/idb/analysis.py b/idb/analysis.py index 7d77485..9e48b0f 100644 --- a/idb/analysis.py +++ b/idb/analysis.py @@ -201,6 +201,13 @@ def addr(self): else: raise RuntimeError('unexpected wordsize') + def off(self): + offset = self.addr() + mask = (2 ** (self.wordsize *8 )) - 1 + if offset & (1 << ((self.wordsize * 8) - 1)): + return offset | ~mask + else: + return offset Field = namedtuple('Field', ['name', 'tag', 'index', 'cast', 'minver']) # namedtuple default args. @@ -626,8 +633,10 @@ def __init__(self, buf, wordsize): # eg. all of these, if high bit of flags not set. pass else: + # We are in a function tail. Chunks can be above or below the tail + # owner try: - self.owner = self.startEA - u.addr() + self.owner = self.startEA - u.off() self.refqty = u.dd() except IndexError: # see warning note above diff --git a/idb/idapython.py b/idb/idapython.py index 2d5e2e2..fee542f 100644 --- a/idb/idapython.py +++ b/idb/idapython.py @@ -1,5 +1,7 @@ # -*- coding: utf-8 -*- +import os import re +import struct import logging import weakref import collections @@ -36,6 +38,35 @@ def cached_method(*args, **kwargs): return decorator +# This decorator is meant to wrap a module in another one like IDA does with +# ida_* wrapped in idaapi and/or idc. +# We could use __getattr__ in idaapi to act as a proxy but that will break +# statements like "from idaapi import *" or "from idc import *", a quite common +# pattern in IDAPython scripts. +# We do it here instead of shim.py:HookedImporter to get the same wraps in +# situations where the shim is not needed +# Use it on ida_* __init__() +# XXX: This is mostly a prototype. Needs further considerations. +def wrap_module(into, full=True): + def decorator(func): + @functools.wraps(func) + def wrapped_func(self, *args, **kwargs): + func(self, *args, **kwargs) + mod = self.api.__dict__[into] + for attr in dir(self): + # Do not set private fields and api/idb objs already present in + # every module + if attr.startswith("_") or attr == 'api' or attr == 'idb': + continue + obj = getattr(self, attr) + # If full is False, only set "constants" + if not full and callable(obj): + continue + setattr(mod, attr, obj) + return wrapped_func + return decorator + + def is_flag_set(flags, flag): return flags & flag == flag @@ -372,6 +403,7 @@ class AFLAGS: class ida_netnode: + @wrap_module("idaapi") def __init__(self, db, api): self.idb = db self.api = api @@ -381,6 +413,8 @@ def netnode(self, *args, **kwargs): class ida_ida: + @wrap_module("idaapi") + @wrap_module("idc", full=False) def __init__(self, db, api): self.idb = db self.api = api @@ -568,6 +602,30 @@ def __init__(self, db, api): self.IDB_EXT64 = 'i64' self.IDB_EXT = 'idb' +class ida_ua: + # op_t + # via: https://www.hex-rays.com/products/ida/support/sdkdoc/group__o__.html + + o_void = 0 # No operand + o_reg = 1 # General register + o_mem = 2 # Direct memory reference + o_phrase = 3 # Memory reference using registers + o_displ = 4 # Memory reference using registers and displacement + o_imm = 5 # Immediate value + o_far = 6 # Immediate far address + o_near = 7 # Immediate near address + o_idpspec0 = 8 # Processor specific + o_idpspec1 = 9 + o_idpspec2 = 10 + o_idpspec3 = 11 + o_idpspec4 = 12 + o_idpspec5 = 13 + + @wrap_module("idaapi") + @wrap_module("idc", full=False) + def __init__(self, db, api): + self.idb = db + self.api = api class idc: @@ -660,6 +718,7 @@ def __init__(self, db, api): self.SEGATTR_COLOR = 100 self.BADADDR = 0xFFFFFFFF + self.__EA64__ = False elif self.idb.wordsize == 8: self.FUNCATTR_START = 0 @@ -693,9 +752,20 @@ def __init__(self, db, api): self.SEGATTR_COLOR = 188 self.BADADDR = 0xFFFFFFFFFFFFFFFF + self.__EA64__ = True else: raise RuntimeError('unexpected wordsize') + # Command line arguments passed to idapython scripts. Args passed via -S + # switch in IDA + self.ARGV = [] + + ## Mantain API compatibility for API < 7 + self.GetMnem = self.print_insn_mnem + self.GetOpnd = self.print_operand + self.GetOpType = self.get_operand_type + self.FindFuncEnd = self.find_func_end + def ScreenEA(self): return self.api.ScreenEA @@ -796,18 +866,7 @@ def Head(self, ea): return ea def ItemSize(self, ea): - oea = ea - flags = self.GetFlags(ea) - if not self.api.ida_bytes.is_head(flags): - raise ValueError('ItemSize must only be called on a head address.') - - ea += 1 - flags = self.GetFlags(ea) - while flags is not None and not self.api.ida_bytes.is_head(flags): - ea += 1 - # TODO: handle Index/KeyError here when we overrun a segment - flags = self.GetFlags(ea) - return ea - oea + return self.api.ida_bytes.get_item_end(ea) - ea def NextHead(self, ea): ea += 1 @@ -901,6 +960,22 @@ def _disassemble(self, ea): dis = self._load_dis(capstone.CS_ARCH_MIPS, capstone.CS_MODE_MIPS32 | capstone.CS_MODE_LITTLE_ENDIAN) elif bitness == 64: dis = self._load_dis(capstone.CS_ARCH_MIPS, capstone.CS_MODE_MIPS64 | capstone.CS_MODE_LITTLE_ENDIAN) + elif procname == "ppc": + if bitness == 32: + dis = self._load_dis(capstone.CS_ARCH_PPC, capstone.CS_MODE_32 | capstone.CS_MODE_BIG_ENDIAN) + elif bitness == 64: + dis = self._load_dis(capstone.CS_ARCH_PPC, capstone.CS_MODE_64 | capstone.CS_MODE_BIG_ENDIAN) + elif procname == "ppcl": + if bitness == 32: + dis = self._load_dis(capstone.CS_ARCH_PPC, capstone.CS_MODE_32 | capstone.CS_MODE_LITTLE_ENDIAN) + elif bitness == 64: + dis = self._load_dis(capstone.CS_ARCH_PPC, capstone.CS_MODE_64 | capstone.CS_MODE_LITTLE_ENDIAN) + elif procname == "sparcb": + if bitness == 32: + dis = self._load_dis(capstone.CS_ARCH_SPARC, capstone.CS_MODE_BIG_ENDIAN) + elif procname == "sparcl": + if bitness == 32: + dis = self._load_dis(capstone.CS_ARCH_SPARC, capstone.CS_MODE_LITTLE_ENDIAN) if dis is None: raise NotImplementedError("unknown arch %s bit:%s inst_len:%d" % (procname, bitness, len(inst_buf))) @@ -913,7 +988,7 @@ def _disassemble(self, ea): else: return op - def GetMnem(self, ea): + def print_insn_mnem(self, ea): op = self._disassemble(ea) return op.mnemonic @@ -921,6 +996,54 @@ def GetDisasm(self, ea): op = self._disassemble(ea) return '%s\t%s' % (op.mnemonic, op.op_str) + def print_operand(self, ea, n): + op = self._disassemble(ea) + opnds = op.op_str.split(", ") + n_opnds = len(opnds) + + if 0 <= n < n_opnds: + return opnds[n] + else: + return "" + + def get_operand_type(self, ea, n): + from capstone import CS_OP_INVALID, CS_OP_REG, CS_OP_MEM, CS_OP_IMM + + op = self._disassemble(ea) + opnds = op.operands + n_opnds = len(opnds) + + # capstone produces 2 operands for immediate far jmp/call, IDA only 1 + # TODO: we need better handling of o_far recognition + is_far = False + if op.mnemonic in ["ljmp", "lcall"]: + n_opnds = 1 + is_far = True + # continue normal operand type check + if 0 <= n < n_opnds: + op_n = opnds[n] + if op_n.type == CS_OP_INVALID: + return -1 + elif op_n.type == CS_OP_REG: + return self.api.ida_ua.o_reg + elif op_n.type == CS_OP_MEM: + op_mem = op_n.value.mem + if op_mem.base == 0: + return self.api.ida_ua.o_mem + if op_mem.base != 0 and op_mem.disp == 0: + return self.api.ida_ua.o_phrase + if op_mem.base != 0 and op_mem.disp != 0: + return self.api.ida_ua.o_displ + elif op_n.type == CS_OP_IMM: + if is_far: + return self.api.ida_ua.o_far + elif self.api.ida_bytes.is_code(self.GetFlags(op_n.value.imm)): + return self.api.ida_ua.o_near + else: + return self.api.ida_ua.o_imm + else: + return self.api.ida_ua.o_void + # one instruction or data CIC_ITEM = 1 # function @@ -985,6 +1108,13 @@ def GetFunctionAttr(self, ea, attr): def GetFunctionName(self, ea): return self.api.ida_funcs.get_func_name(ea) + def find_func_end(self, ea): + func = self.api.ida_funcs.get_func(ea) + if not func: + return self.BADADDR + else: + return func.endEA + def LocByName(self, name): try: key = ("N" + name).encode('utf-8') @@ -1000,7 +1130,7 @@ def GetInputSHA256(self): return self.api.ida_nalt.retrieve_input_file_sha256() def GetInputFile(self): - return self.api.ida_nalt.get_input_file_path() + return os.path.basename(self.api.ida_nalt.get_input_file_path()) def Comment(self, ea): return self.api.ida_bytes.get_cmt(ea, False) @@ -1153,10 +1283,14 @@ def LineB(self, ea, num): class ida_bytes: + @wrap_module("idaapi") def __init__(self, db, api): self.idb = db self.api = api + ## Mantain API compatibility for API < 7 + self.get_long = self.get_dword + def get_cmt(self, ea, repeatable): flags = self.api.idc.GetFlags(ea) if not self.has_cmt(flags): @@ -1338,8 +1472,40 @@ def next_not_tail(self, ea): def next_inited(self, ea, maxea): return self.next_that(ea, maxea, lambda flags: ida_bytes.has_value(flags)) + def get_item_end(self, ea): + ea += 1 + flags = self.api.idc.GetFlags(ea) + while flags is not None and not self.api.ida_bytes.is_head(flags) and self.api.idc.SegEnd(ea): + ea += 1 + flags = self.api.idc.GetFlags(ea) + return ea + + def get_byte(self, ea): + return ord(self.get_bytes(ea, 1)) + + def get_word(self, ea): + if self.api.idaapi.get_inf_structure().is_be: + fmt = ">H" + else: + fmt = "