Python graph.DiGraph類代碼示例

def gen_bloc_data_flow_graph(ir_arch, ad, block_flow_cb):
    for irbloc in ir_arch.blocs.values():
        print irbloc

    ir_arch.gen_graph()
    ir_arch.dead_simp()

    irbloc_0 = None
    for irbloc in ir_arch.blocs.values():
        if irbloc.label.offset == ad:
            irbloc_0 = irbloc
            break
    assert(irbloc_0 is not None)
    flow_graph = DiGraph()
    flow_graph.node2str = lambda n: node2str(flow_graph, n)

    for irbloc in ir_arch.blocs.values():
        block_flow_cb(ir_arch, flow_graph, irbloc)

    for irbloc in ir_arch.blocs.values():
        print irbloc
        print 'IN', [str(x) for x in irbloc.in_nodes]
        print 'OUT', [str(x) for x in irbloc.out_nodes]

    print '*' * 20, 'interbloc', '*' * 20
    inter_bloc_flow(ir_arch, flow_graph, irbloc_0.label)

    # from graph_qt import graph_qt
    # graph_qt(flow_graph)
    open('data.dot', 'w').write(flow_graph.dot())

class basicblocs:

    def __init__(self, ab=[]):
        self.blocs = {}
        self.g = DiGraph()
        self.add_blocs(ab)

    def add(self, b):
        self.blocs[b.label] = b
        self.g.add_node(b.label)
        for dst in b.bto:
            if isinstance(dst.label, asm_label):
                self.g.add_edge(b.label, dst.label)

    def add_blocs(self, ab):
        for b in ab:
            self.add(b)

    def get_bad_dst(self):
        o = set()
        for b in self.blocs.values():
            for c in b.bto:
                if c.c_t == asm_constraint.c_bad:
                    o.add(b)
        return o

def unflatGraph(flat_graph):
    graph = DiGraph()
    nodes, edges = flat_graph
    for node in nodes:
        graph.add_node(node)
    for nodeA, nodeB in edges:
        graph.add_edge(nodeA, nodeB)
    return graph

    def as_graph(self, starting_nodes):
        """Return a DiGraph corresponding to computed dependencies, with
        @starting_nodes as leafs
        @starting_nodes: set of DependencyNode instance
        """

        # Build subgraph for each starting_node
        subgraphs = []
        for starting_node in starting_nodes:
            subgraphs.append(self._build_depGraph(starting_node))

        # Merge subgraphs into a final DiGraph
        graph = DiGraph()
        for sourcegraph in subgraphs:
            for node in sourcegraph.nodes():
                graph.add_node(node)
            for edge in sourcegraph.edges():
                graph.add_uniq_edge(*edge)
        return graph

def gen_bloc_data_flow_graph(ir_arch, in_str, ad):  # arch, attrib, pool_bin, bloc, symbol_pool):
    out_str = ""

    # ir_arch = ir_x86_32(symbol_pool)

    for irbloc in ir_arch.blocs.values():
        print irbloc

    ir_arch.gen_graph()
    ir_arch.dead_simp()

    irbloc_0 = None
    for irbloc in ir_arch.blocs.values():
        if irbloc.label.offset == ad:
            irbloc_0 = irbloc
            break
    assert(irbloc_0 is not None)
    flow_graph = DiGraph()
    flow_graph.node2str = lambda n: node2str(flow_graph, n)
    done = set()
    todo = set([irbloc_0.label])

    bloc2w = {}

    for irbloc in ir_arch.blocs.values():
        intra_bloc_flow_raw(ir_arch, flow_graph, irbloc)
        # intra_bloc_flow_symb(ir_arch, flow_graph, irbloc)

    for irbloc in ir_arch.blocs.values():
        print irbloc
        print 'IN', [str(x) for x in irbloc.in_nodes]
        print 'OUT', [str(x) for x in irbloc.out_nodes]

    print '*' * 20, 'interbloc', '*' * 20
    inter_bloc_flow(ir_arch, flow_graph, irbloc_0.label)

    # sys.path.append('/home/serpilliere/projet/m2_devel/miasm2/core')
    # from graph_qt import graph_qt
    # graph_qt(flow_graph)
    open('data.txt', 'w').write(flow_graph.dot())

def gen_block_data_flow_graph(ir_arch, ircfg, ad, block_flow_cb):
    for irblock in ircfg.blocks.values():
        print irblock

    dead_simp(ir_arch, ircfg)


    irblock_0 = None
    for irblock in ircfg.blocks.values():
        loc_key = irblock.loc_key
        offset = ircfg.loc_db.get_location_offset(loc_key)
        if offset == ad:
            irblock_0 = irblock
            break
    assert(irblock_0 is not None)
    flow_graph = DiGraph()
    flow_graph.node2str = lambda n: node2str(flow_graph, n)


    irb_in_nodes = {}
    irb_out_nodes = {}
    for label in ircfg.blocks:
        irb_in_nodes[label] = {}
        irb_out_nodes[label] = {}

    for label, irblock in ircfg.blocks.iteritems():
        block_flow_cb(ir_arch, ircfg, flow_graph, irblock, irb_in_nodes[label], irb_out_nodes[label])

    for label in ircfg.blocks:
        print label
        print 'IN', [str(x) for x in irb_in_nodes[label]]
        print 'OUT', [str(x) for x in irb_out_nodes[label]]

    print '*' * 20, 'interblock', '*' * 20
    inter_block_flow(ir_arch, ircfg, flow_graph, irblock_0.loc_key, irb_in_nodes, irb_out_nodes)

    # from graph_qt import graph_qt
    # graph_qt(flow_graph)
    open('data.dot', 'w').write(flow_graph.dot())

 def as_graph(self):
     """Generates a Digraph of dependencies"""
     graph = DiGraph()
     for node_a, node_b in self.links:
         if not node_b:
             graph.add_node(node_a)
         else:
             graph.add_edge(node_a, node_b)
     for parent, sons in self.pending.iteritems():
         for son in sons:
             graph.add_edge(parent, son)
     return graph

 def gen_graph(self, link_all=True):
     """
     Gen irbloc digraph
     @link_all: also gen edges to non present irblocs
     """
     self.g = DiGraph()
     for lbl, b in self.blocs.items():
         # print 'add', lbl
         self.g.add_node(lbl)
         # dst = self.get_bloc_dst(b)
         dst = self.dst_trackback(b)
         # print "\tdst", dst
         for d in dst:
             if isinstance(d, ExprInt):
                 d = ExprId(self.symbol_pool.getby_offset_create(int(d.arg)))
             if self.ExprIsLabel(d):
                 if d.name in self.blocs or link_all is True:
                     self.g.add_edge(lbl, d.name)

def blist2graph(ab):
    """
    ab: list of asmbloc
    return: graph of asmbloc
    """
    g = DiGraph()
    g.lbl2bloc = {}
    for b in ab:
        g.lbl2bloc[b.label] = b
        g.add_node(b.label)
        for x in b.bto:
            g.add_edge(b.label, x.label)
    return g

    def __init__(self, abicls, machine):
        self.abicls = abicls

        self.input_reg = {}
        self.output_reg = {}

        self._previous_addr = 0
        self._current_addr = 0
        self._instr_count = 0
        self._pending_call = []
        # Function addr -> list of information on calls
        self.function_calls = {}
        self.paths = DiGraph()

        self.in_memory = {}
        self.out_memory = {}

        self._ira = Machine(machine).ira()
        self._ptr_size = self._ira.sizeof_pointer()/8
        self.sp = self._ira.sp.name

    def _build_depGraph(self, depnode):
        """Recursively build the final list of DiGraph, and clean up unmodifier
        nodes
        @depnode: starting node
        """

        if depnode not in self._cache or \
                not self._cache[depnode]:
            ## There is no dependency
            graph = DiGraph()
            graph.add_node(depnode)
            return graph

        # Recursion
        dependencies = list(self._cache[depnode])

        graphs = []
        for sub_depnode in dependencies:
            graphs.append(self._build_depGraph(sub_depnode))

        # head(graphs[i]) == dependencies[i]
        graph = DiGraph()
        graph.add_node(depnode)
        for head in dependencies:
            graph.add_uniq_edge(head, depnode)

        for subgraphs in itertools.product(graphs):
            for sourcegraph in subgraphs:
                for node in sourcegraph.nodes():
                    graph.add_node(node)
                for edge in sourcegraph.edges():
                    graph.add_uniq_edge(*edge)

        # Update the running queue
        return graph

class Snapshot(object):

    @classmethod
    def get_byte(cls, value, byte):
        '''Return the byte @byte of the value'''
        return struct.pack('@B', (value & (0xFF << (8 * byte))) >> (8 * byte))

    @classmethod
    def unpack_ptr(cls, value):
        return struct.unpack('@P', value)[0]

    def __init__(self, abicls, machine):
        self.abicls = abicls

        self.input_reg = {}
        self.output_reg = {}

        self._previous_addr = 0
        self._current_addr = 0
        self._instr_count = 0
        self._pending_call = []
        # Function addr -> list of information on calls
        self.function_calls = {}
        self.paths = DiGraph()

        self.in_memory = {}
        self.out_memory = {}

        self._ira = Machine(machine).ira()
        self._ptr_size = self._ira.sizeof_pointer()/8
        self.sp = self._ira.sp.name

    def add_input_register(self, reg_name, reg_value):
        self.input_reg[reg_name] = reg_value

    def add_output_register(self, reg_name, reg_value):
        self.output_reg[reg_name] = reg_value

    def add_memory_read(self, address, size, value):
        for i in xrange(size):
            self.out_memory[address + i] = MemoryAccess(1,
                                                        Snapshot.get_byte(value, i),
                                                        0,  # Output access never used
            )

            if address + i not in self.in_memory:
                self.in_memory[address + i] = MemoryAccess(1,
                                                           Snapshot.get_byte(value, i),
                                                           PAGE_READ,
                )

            else:
                self.in_memory[address + i].access |= PAGE_READ

    def add_memory_write(self, address, size, value):
        for i in xrange(size):
            self.out_memory[address + i] = MemoryAccess(1,
                                                        Snapshot.get_byte(value, i),
                                                        0,  # Output access never used
            )

            if address + i not in self.in_memory:
                self.in_memory[address + i] = MemoryAccess(1,
                                                           "\x00",
                                                           # The value is
                                                           # not used by the
                                                           # test
                                                           PAGE_WRITE,
                )

            else:
                self.in_memory[address + i].access |= PAGE_WRITE

    def add_executed_instruction(self, address):
        '''
        Function called to signal that the address has been executed
        This function has to be called in the order of their executed instruction
        Else paths can not be updated correctly
        '''
        self._previous_addr = self._current_addr
        self._current_addr = address
        self.paths.add_uniq_edge(self._previous_addr, self._current_addr)
        self._instr_count += 1

        # Resolve call destination
        if (self._pending_call and
            self._previous_addr == self._pending_call[-1]["caller_addr"]):
            info = self._pending_call[-1]
            info["dest"] = address
            info["beg"] = self._instr_count


    def add_call(self, caller_addr, stack_ptr):
        '''
        Function call, target is not determined yet
        called *before* instruction execution
        '''
        info = {"stack_ptr": stack_ptr,
                "caller_addr": caller_addr,
        }
#.........這裏部分代碼省略.........

class ira:

    def sort_dst(self, todo, done):
        out = set()
        while todo:
            dst = todo.pop()
            if self.ExprIsLabel(dst):
                done.add(dst)
            elif isinstance(dst, ExprMem) or isinstance(dst, ExprInt):
                done.add(dst)
            elif isinstance(dst, ExprCond):
                todo.add(dst.src1)
                todo.add(dst.src2)
            elif isinstance(dst, ExprId):
                out.add(dst)
            else:
                done.add(dst)
        return out

    def dst_trackback(self, b):
        dst = b.dst
        todo = set([dst])
        out = set()
        done = set()

        for irs in reversed(b.irs):
            if len(todo) == 0:
                break
            out = self.sort_dst(todo, done)
            found = set()
            follow = set()
            for i in irs:
                if not out:
                    break
                for o in out:
                    if i.dst == o:
                        follow.add(i.src)
                        found.add(o)
                for o in found:
                    out.remove(o)

            for o in out:
                if not o in found:
                    follow.add(o)
            todo = follow
        out = self.sort_dst(todo, done)

        return done

    def gen_graph(self, link_all = False):
        """
        Gen irbloc digraph
        @link_all: also gen edges to non present irblocs
        """
        self.g = DiGraph()
        for lbl, b in self.blocs.items():
            # print 'add', lbl
            self.g.add_node(lbl)
            # dst = self.get_bloc_dst(b)
            dst = self.dst_trackback(b)
            # print "\tdst", dst
            for d in dst:
                if isinstance(d, ExprInt):
                    d = ExprId(
                        self.symbol_pool.getby_offset_create(int(d.arg)))
                if self.ExprIsLabel(d):
                    if d.name in self.blocs or link_all is True:
                        self.g.add_edge(lbl, d.name)

    def graph(self):
        out = """
    digraph asm_graph {
    size="80,50";
    node [
    fontsize = "16",
    shape = "box"
    ];
    """
        all_lbls = {}
        for lbl in self.g.nodes():
            if not lbl in self.blocs:
                continue
            b = self.blocs[lbl]
            ir_txt = [str(lbl)]
            for irs in b.irs:
                for l in irs:
                    ir_txt.append(str(l))
                ir_txt.append("")
            ir_txt.append("")
            all_lbls[id(lbl)] = "\l\\\n".join(ir_txt)
        for l, v in all_lbls.items():
            out += '%s [label="%s"];\n' % (l, v)

        for a, b in self.g.edges():
            out += '%s -> %s;\n' % (id(a), id(b))
        out += '}'
        return out

    def remove_dead(self, b):
        for ir, _, c_out in zip(b.irs, b.c_in, b.c_out):
#.........這裏部分代碼省略.........

class ira:
    def ira_regs_ids(self):
        """Returns ids of all registers used in the IR"""
        return self.arch.regs.all_regs_ids + [self.IRDst]

    def sort_dst(self, todo, done):
        out = set()
        while todo:
            dst = todo.pop()
            if self.ExprIsLabel(dst):
                done.add(dst)
            elif isinstance(dst, ExprMem) or isinstance(dst, ExprInt):
                done.add(dst)
            elif isinstance(dst, ExprCond):
                todo.add(dst.src1)
                todo.add(dst.src2)
            elif isinstance(dst, ExprId):
                out.add(dst)
            else:
                done.add(dst)
        return out

    def dst_trackback(self, b):
        dst = b.dst
        todo = set([dst])
        done = set()

        for irs in reversed(b.irs):
            if len(todo) == 0:
                break
            out = self.sort_dst(todo, done)
            found = set()
            follow = set()
            for i in irs:
                if not out:
                    break
                for o in out:
                    if i.dst == o:
                        follow.add(i.src)
                        found.add(o)
                for o in found:
                    out.remove(o)

            for o in out:
                if o not in found:
                    follow.add(o)
            todo = follow

        return done

    def gen_graph(self, link_all=True):
        """
        Gen irbloc digraph
        @link_all: also gen edges to non present irblocs
        """
        self.g = DiGraph()
        for lbl, b in self.blocs.items():
            # print 'add', lbl
            self.g.add_node(lbl)
            # dst = self.get_bloc_dst(b)
            dst = self.dst_trackback(b)
            # print "\tdst", dst
            for d in dst:
                if isinstance(d, ExprInt):
                    d = ExprId(self.symbol_pool.getby_offset_create(int(d.arg)))
                if self.ExprIsLabel(d):
                    if d.name in self.blocs or link_all is True:
                        self.g.add_edge(lbl, d.name)

    def graph(self):
        """Output the graphviz script"""
        out = """
    digraph asm_graph {
    size="80,50";
    node [
    fontsize = "16",
    shape = "box"
    ];
        """
        all_lbls = {}
        for lbl in self.g.nodes():
            if lbl not in self.blocs:
                continue
            irb = self.blocs[lbl]
            ir_txt = [str(lbl)]
            for irs in irb.irs:
                for l in irs:
                    ir_txt.append(str(l))
                ir_txt.append("")
            ir_txt.append("")
            all_lbls[hash(lbl)] = "\l\\\n".join(ir_txt)
        for l, v in all_lbls.items():
            # print l, v
            out += '%s [label="%s"];\n' % (l, v)

        for a, b in self.g.edges():
            # print 'edge', a, b, hash(a), hash(b)
            out += "%s -> %s;\n" % (hash(a), hash(b))
        out += "}"
        return out
#.........這裏部分代碼省略.........

class ira:

    def sort_dst(self, todo, done):
        out = set()
        while todo:
            dst = todo.pop()
            if self.ExprIsLabel(dst):
                done.add(dst)
            elif isinstance(dst, ExprMem) or isinstance(dst, ExprInt):
                done.add(dst)
            elif isinstance(dst, ExprCond):
                todo.add(dst.src1)
                todo.add(dst.src2)
            elif isinstance(dst, ExprId):
                out.add(dst)
            else:
                done.add(dst)
        return out

    def dst_trackback(self, b):
        dst = b.dst
        todo = set([dst])
        out = set()
        done = set()

        for irs in reversed(b.irs):
            if len(todo) == 0:
                break
            out = self.sort_dst(todo, done)
            found = set()
            follow = set()
            for i in irs:
                if not out:
                    break
                for o in out:
                    if i.dst == o:
                        follow.add(i.src)
                        found.add(o)
                for o in found:
                    out.remove(o)

            for o in out:
                if o not in found:
                    follow.add(o)
            todo = follow
        out = self.sort_dst(todo, done)

        return done

    def gen_graph(self, link_all = True):
        """
        Gen irbloc digraph
        @link_all: also gen edges to non present irblocs
        """
        self.g = DiGraph()
        for lbl, b in self.blocs.items():
            # print 'add', lbl
            self.g.add_node(lbl)
            # dst = self.get_bloc_dst(b)
            dst = self.dst_trackback(b)
            # print "\tdst", dst
            for d in dst:
                if isinstance(d, ExprInt):
                    d = ExprId(
                        self.symbol_pool.getby_offset_create(int(d.arg)))
                if self.ExprIsLabel(d):
                    if d.name in self.blocs or link_all is True:
                        self.g.add_edge(lbl, d.name)

    def graph(self):
        """Output the graphviz script"""
        out = """
    digraph asm_graph {
    size="80,50";
    node [
    fontsize = "16",
    shape = "box"
    ];
        """
        all_lbls = {}
        for lbl in self.g.nodes():
            if lbl not in self.blocs:
                continue
            irb = self.blocs[lbl]
            ir_txt = [str(lbl)]
            for irs in irb.irs:
                for l in irs:
                    ir_txt.append(str(l))
                ir_txt.append("")
            ir_txt.append("")
            all_lbls[hash(lbl)] = "\l\\\n".join(ir_txt)
        for l, v in all_lbls.items():
            # print l, v
            out += '%s [label="%s"];\n' % (l, v)

        for a, b in self.g.edges():
            # print 'edge', a, b, hash(a), hash(b)
            out += '%s -> %s;\n' % (hash(a), hash(b))
        out += '}'
        return out
#.........這裏部分代碼省略.........