Skip to content

Examples | Creating a dotfile of tools in your home directory

Jump to bottom Edit New page
Ali Rizvi-Santiago edited this page Nov 18, 2022 · 2 revisions

Ideally, you'd drop these in your $HOME/.idapythonrc.py file or $YOURIDBPATH/idapythonrc.py so that way you have access to them when you start things up.

Random graph tools

import networkx as nx
def complexity(G):
    edges = len(G.edges())
    nodes = len(G.nodes())
    parts = nx.components.number_strongly_connected_components(G)
    return edges - nodes + parts

def loops(G):
    for blocks in sorted(nx.simple_cycles(G), key=len):
        yield sorted(blocks)
    return

def paths(G):
    g = G.to_undirected()
    for blocks in sorted(nx.cycle_basis(g)):
        yield sorted(blocks)
    return

def path(G, start, stop):
    f = func.by(G.name)
    b1, b2 = func.block(start), func.block(stop)
    items = nx.shortest_path(G, func.block.address(b1), func.block.address(b2))
    return [func.block(item) for item in items]

Savng or loading arbitrary data in the same directory as the database

import bz2, pickle
def load(file):
    with open(file if file.startswith('/') else db.config.path(file), 'rb') as infile:
        data = infile.read()
    return pickle.loads(bz2.decompress(data))

def save(file, data):
    compressed = bz2.compress(pickle.dumps(data))
    with open(file if file.startswith('/') else db.config.path(file), 'xb') as outfile:
        outfile.write(compressed)
    return

Simple cfuncptr_t tools

def find_expr_addr(cfunc, cexpr):
    if cexpr.ea == idaapi.BADADDR:
        while True:
            cexpr = cfunc.body.find_parent_of(cexpr)
            if cexpr is None:
                ea = cfunc.entry_ea
                break
            if cexpr.ea != idaapi.BADADDR:
                ea = cexpr.ea
                break
            continue
        return ea
    return cexpr.ea

def find_addr(vu):
    citem = vu.item
    if citem.citype in {idaapi.VDI_EXPR}:
        return find_expr_addr(vu.cfunc, citem.e)
    elif citem.citype in {idaapi.VDI_TAIL}:
        return citem.loc.ea
    elif citem.citype in {idaapi.VDI_LVAR}:
        return citem.l.defea
    elif citem.citype in {idaapi.VDI_FUNC}:
        return citem.f.entry_ea
    raise NotImplementedError(citem.citype)

import internal
class lvars(object):
    @classmethod
    def __iterate__(cls, D):
        lvars = D.get_lvars()
        for index in range(lvars.size()):
            yield lvars[index]
        return

    __matcher__ = internal.utils.matcher()
    __matcher__.boolean('name', lambda name, item: name.lower() == item.lower(), fgetattr('name'))
    __matcher__.combinator('like', utils.fcompose(fnmatch.translate, utils.fpartial(re.compile, flags=re.IGNORECASE), operator.attrgetter('match')), fgetattr('name'))
    __matcher__.predicate('predicate'), __matcher__.predicate('pred')

    @internal.utils.multicase()
    @classmethod
    def iterate(cls, **type):
        return cls.iterate(None, **type)
    @internal.utils.multicase(name=str)
    @classmethod
    def iterate(cls, name, **type):
        return cls.iterate(name, None, **type)
    @internal.utils.multicase(name=str, D=(None.__class__, idaapi.cfuncptr_t))
    @classmethod
    def iterate(cls, name, D, **type):
        type.setdefault('like', name)
        return cls.iterate(D, **type)
    @internal.utils.multicase(D=(None.__class__, idaapi.cfuncptr_t))
    @classmethod
    def iterate(cls, D, **type):
        state = D if D else idaapi.decompile(func.address())
        iterable = cls.__iterate__(state)
        for key, value in (type or {'predicate': utils.fconstant(True)}).items():
            iterable = cls.__matcher__.match(key, value, iterable)
        for item in iterable: yield item

    @internal.utils.multicase()
    @classmethod
    def list(cls, **type):
        return cls.list(None, **type)
    @internal.utils.multicase(name=str)
    @classmethod
    def list(cls, name, **type):
        return cls.list(name, None, **type)
    @internal.utils.multicase(name=str, D=(None.__class__, idaapi.cfuncptr_t))
    @classmethod
    def list(cls, name, D, **type):
        type.setdefault('like', name)
        return cls.list(D, **type)
    @internal.utils.multicase(D=(None.__class__, idaapi.cfuncptr_t))
    @classmethod
    def list(cls, D, **type):
        state = D if D else idaapi.decompile(func.address())
        res = [item for item in cls.iterate(state, **type)]

        print_t = lambda item: idaapi.print_tinfo('', 0, 0, 0, item.type(), item.name, '')

        maxdisasm = max(map(fcompose(fgetattr('defea'), db.disasm, len), res))
        maxname = max(map(fcompose(print_t, len), res))

        for item in res:
            t_s = print_t(item)
            print("{:<{:d}s} // {:<{:d}s} : {!s}".format(db.disasm(item.defea), maxdisasm, t_s, maxname, cls.vdloc(state, item)))
        return

    @internal.utils.multicase()
    @classmethod
    def by(cls, **type):
        return cls.by(None, **type)
    @internal.utils.multicase(name=str)
    @classmethod
    def by(cls, name, **type):
        return cls.by(name, None, **type)
    @internal.utils.multicase(name=str, D=(None.__class__, idaapi.cfuncptr_t))
    @classmethod
    def by(cls, name, D, **type):
        type.setdefault('like', name)
        return cls.by(D, **type)
    @internal.utils.multicase(D=(None.__class__, idaapi.cfuncptr_t))
    @classmethod
    def by(cls, D, **type):
        state = D if D else idaapi.decompile(func.address())
        return next(item for item in cls.iterate(state, **type))

    @internal.utils.multicase()
    @classmethod
    def get(cls, **type):
        return cls.get(None, **type)
    @internal.utils.multicase(name=str)
    @classmethod
    def get(cls, name, **type):
        return cls.get(name, None, **type)
    @internal.utils.multicase(name=str, D=(None.__class__, idaapi.cfuncptr_t))
    @classmethod
    def get(cls, name, D, **type):
        type.setdefault('like', name)
        return cls.get(D, **type)
    @internal.utils.multicase(D=(None.__class__, idaapi.cfuncptr_t))
    @classmethod
    def get(cls, D, **type):
        state = D if D else idaapi.decompile(func.address())
        res = cls.by(state, **type)
        return cls.vdloc(state, res)

    @internal.utils.multicase()
    @classmethod
    def name(cls, **type):
        return cls.name(None, **type)
    @internal.utils.multicase(name=str)
    @classmethod
    def name(cls, name, **type):
        return cls.name(name, None, **type)
    @internal.utils.multicase(name=str, D=(None.__class__, idaapi.cfuncptr_t))
    @classmethod
    def name(cls, name, D, **type):
        type.setdefault('like', name)
        return cls.name(D, **type)
    @internal.utils.multicase(D=(None.__class__, idaapi.cfuncptr_t))
    @classmethod
    def name(cls, D, **type):
        state = D if D else idaapi.decompile(func.address())
        res = cls.by(state, **type)
        return res.name

    @classmethod
    def vdloc(cls, state, lv):
        loc = lv.location
        atype, vtype = loc.atype(), lv.type()
        if atype in {idaapi.ALOC_REG1}:
            regname = idaapi.print_vdloc(loc, vtype.get_size())
            return ins.arch.byname(regname)
        elif atype in {idaapi.ALOC_STACK}:
            fr = func.frame(lv.defea)
            delta = state.get_stkoff_delta()
            realoffset = loc.stkoff() - delta
            return fr.members.by_realoffset(realoffset) if 0 <= realoffset < fr.size else location_t(realoffset, vtype.get_size())
        raise NotImplementedError(atype)

    @classmethod
    def refs(cls, lv):
        raise NotImplementedError('python/examples/hexrays/vds12.py')
Clone this wiki locally