trace_path2.py

import sys
import pprint
from copy import copy
from itertools import tee, izip, dropwhile, imap, groupby, ifilterfalse
from operator import itemgetter
from netaddr import IPAddress, IPNetwork
from trigger.cmds import Commando
from trigger.netdevices import NetDevices
from collections import defaultdict, namedtuple
from graphviz import Digraph
from trigger.acl import parse
from twisted.python import log
#log.startLogging(sys.stdout, setStdout=False)


class ShowIpInterface(Commando):
    vendors = ['cisco']
    #commands = ['show ip int brief']
    commands = ['show ip route', 'show interfaces']


class Nodes(dict):
    def add(self, node):
        assert(isinstance(node, Node))
        self.__setitem__(str(node), node)

    def find(self, node):
        return self.get(node, None)

    def all(self):
        return self.values()

    def who_has_route(self, route, routing_table):
        return self.find(str(filter(lambda x: x.protocol == 'C', sorted(d[route], key=lambda x: x.metric)).pop().hostname))

    def trace(self, src_ip, dst_ip, routing_table):
        src = self.who_has_route(ip_addr_src, routing_table)
        dst = self.who_has_route(ip_addr_dst, routing_table)
        
        it = self._trace_next(None, dst_ip, routing_table)
        it.next()
        paths = it.send(src)
        it.next()
	done = False
	seen = [src]
        idx = 0
	fork = None

	if paths[0] == dst:
	    yield str(src), str(dst)
	    raise StopIteration

        while paths and idx < len(routing_table):
            if not fork and len(paths) > 1:
                fork = seen[-1]
            path = paths.pop(0)
            seen.append(path)
            if dst == path and not fork:
	        yield str(seen[-2]), str(path)
		break
	    elif fork:
	        if seen.count(path) > 1:
	            times_seen = seen.count(path)
		    extract = seen[seen.index(fork):-times_seen]
		    extract.append(path)
		    for i in ((extract[0], i) for i in extract[1:-1]):
		        yield str(i[0]), str(i[1])
		    for i in ((i, extract[-1]) for i in extract[1:-1]):
		        yield str(i[0]), str(i[1])
		    fork = None
		    continue
            else:
                yield str(seen[-2]), str(path)
            p = it.send(path)
            it.next()
	    paths.extend(p)
	    idx += 1


    def _trace_next(self, src, dst_ip, routing_table):
        safety = len(routing_table)
        done = False
        while not done and safety:
	    safety -= 1
	    if src:
	        new_srcs = []
	        for outif in src.get_route(dst_ip):
                    try:
		        adj = src.adjacency(src.connection(outif)[0])
                        new_src = self.find(adj[0])
	                new_srcs.append(new_src)
                    except:
                        done = True
	        yield new_srcs
	        src = (yield src)
	    else:
	        src = (yield src)
        
        


class Node(object):
    def __init__(self, node, connections=None):
        self._node = node
        self._connections = Connections(connections)
	self._adjacencies = {}
	self._routes = {}

    def __str__(self):
        return str(self._node.nodeName)

    def __repr__(self):
        return "Node({})".format(str(self._node.nodeName))

    def connection(self, intf_name):
        return self._connections.get(intf_name)

    def connections_by_intf(self):
        return self._connections

    def connections_by_ip(self):
        return {v[0]: (k, v[1]) for k,v in self.connections_by_intf().iteritems()}

    def adjacency(self, key):
	try:
	    IPNetwork(key)
	    return self.adjacencies_by_ip().get(key)
	except ValueError:
            return self._adjacencies.get(key, None)

    def add_adjacency(self, key, value):
        self._adjacencies[key] = value
        
    def adjacencies_by_intf(self):
        return self._adjacencies

    def adjacencies_by_ip(self):
        return {v[0]: (k, v[1]) for k,v in self._adjacencies.iteritems() }

    def set_routes(self, routing_table):
        self._routes = routing_table

    def get_routes(self):
        return self._routes

    def get_route(self, route):
        return self._routes.get(route, None)


class Connections(dict):
    def add(self, connection):
        self.update(connection)
        
    def to_ip(self):
        return Connections({v: k for k,v in self.iteritems()})

    def to_subnet(self):
        return Connections({v[1]: (v[0], k) for k,v in self.iteritems() if v[1]})

    def interfaces(self):
        return self.keys()

    def ips(self):
        return self.to_ip().keys()
	    

device_list = ['p1.demo.localdomain', 'p2.demo.localdomain', 'pe1.demo.localdomain', 'pe2.demo.localdomain', 'p3.demo.localdomain', 'p4.demo.localdomain']
showipinterface = ShowIpInterface(devices=device_list)
showipinterface.run() # Commando exposes this to start the event loop

#print '\nResults:'
results = showipinterface.parsed_results
nd = NetDevices()
# Process the structured output of the routing table into something we can traverse.

Segment = namedtuple('Segment', ('A', 'B'))
ints = defaultdict(dict)
ints_by_ip = defaultdict(dict)
for hname, data in results.items():
    data = data['show interfaces']
    [ints[hname].update({i[0]: (i[1], i[1] and IPNetwork(i[1]).cidr) or ''}) for i in zip(data['interface'], data['ip_address']) if i[0]]
    [ints_by_ip[hname].update({IPNetwork(i[0]).cidr: (i[1], i[1] and IPNetwork(i[0]).cidr) or ''}) for i in zip(data['ip_address'], data['interface']) if i[0]]


nodes = Nodes()
[nodes.add(Node(i, connections=ints[i.nodeName])) for i in nd.all()]

d = defaultdict(list)
NextHop = namedtuple('NextHop', ('hostname', 'nexthopip', 'protocol', 'metric', 'nexthopif', 'localip'))
for hname, data in results.items():
    data = data['show ip route']
    node = nodes.find(hname)
    routes = map(lambda x: (x[0], x[1] or [ints_by_ip[hname].get(i) for i in ints_by_ip[hname] if IPNetwork(x[2]) in i].pop()[0]), zip((i[0] + i[1] for i in zip(data['network'], data['mask'])), data['nexthopif'], data['nexthopip']))
    routes_parsed = defaultdict(list)
    [routes_parsed[i[0]].append(i[1]) for i in routes]
    node.set_routes(routes_parsed)

    [d["{}{}".format(i[0], i[1])].extend([NextHop(nd.find(hname), i[2], i[3], i[4], i[5], ints[hname].get(i[5]))]) for i in zip(data['network'], data['mask'], data['nexthopip'], data['protocol'], data['metric'], data['nexthopif'])]



rtr_by_hosts = defaultdict(list)
for k,v in d.iteritems():
    for i in v:
        if i.nexthopip:
            rtr_by_hosts[i.hostname.nodeName].append({k: i})

for node in nodes.all():
    for subnet, intf in node.connections_by_intf().to_subnet().iteritems():
        if d.get(str(subnet)):
            for h in d.get(str(subnet)):
                if h.protocol == 'C'  and str(h.hostname) != str(node):
		    node.add_adjacency(str(nodes.find(str(h.hostname))), intf)

#ip_addr_src = '1.1.1.1/32'
#ip_addr_dst = '4.4.4.4/32'
#ip_addr_src = '1.1.1.1/32'
#ip_addr_dst = '2.2.2.2/32'
ip_addr_src = '1.1.1.1/32'
ip_addr_dst = '20.20.20.20/32'
#ip_addr_src = '10.10.10.10/32'
#ip_addr_dst = '20.20.20.20/32'
#ip_addr_src = '10.10.10.10/32'
#ip_addr_dst = '2.2.2.2/32'
#ip_addr_src = '3.3.3.3/32'
#ip_addr_dst = '20.20.20.20/32'
def add_nodes(graph, nodes):
    for n in nodes:
        if isinstance(n, tuple):
            graph.node(n[0], **n[1])
        else:
            graph.node(n)
    return graph

def add_edges(graph, edges):
    for e in edges:
        if isinstance(e[0], tuple):
            graph.edge(*e[0], **e[1])
        else:
            graph.edge(*e)
    return graph

dot = Digraph(comment="Routing path between {} and {}".format(ip_addr_src, ip_addr_dst))
edges = list([i for i in nodes.trace(ip_addr_src, ip_addr_dst, d)])
add_edges(dot, edges).render('img/g5')