ffmap-backend/lib/graph.py

77 lines
2.6 KiB
Python
Raw Normal View History

from functools import reduce
from itertools import chain
import networkx as nx
from lib.nodes import build_mac_table
2015-03-24 16:49:37 +01:00
def import_vis_data(graph, nodes, vis_data):
2015-03-24 16:49:37 +01:00
macs = build_mac_table(nodes)
nodes_a = map(lambda d: 2*[d['primary']], filter(lambda d: 'primary' in d, vis_data))
nodes_b = map(lambda d: [d['secondary'], d['of']], filter(lambda d: 'secondary' in d, vis_data))
graph.add_nodes_from(map(lambda a, b: (a, dict(primary=b, node_id=macs.get(b))), *zip(*chain(nodes_a, nodes_b))))
edges = filter(lambda d: 'neighbor' in d, vis_data)
graph.add_edges_from(map(lambda d: (d['router'], d['neighbor'], dict(tq=float(d['label']))), edges))
def mark_vpn(graph, vpn_macs):
2015-03-24 16:49:37 +01:00
components = map(frozenset, nx.weakly_connected_components(graph))
components = filter(vpn_macs.intersection, components)
nodes = reduce(lambda a, b: a | b, components, set())
for node in nodes:
for k, v in graph[node].items():
v['vpn'] = True
def to_multigraph(graph):
2015-03-24 16:49:37 +01:00
def f(a):
node = graph.node[a]
return node['primary'] if node else a
2015-03-24 16:49:37 +01:00
def map_node(node, data):
return (data['primary'], dict(node_id=data['node_id'])) if data else (node, dict())
digraph = nx.MultiDiGraph()
digraph.add_nodes_from(map(map_node, *zip(*graph.nodes_iter(data=True))))
digraph.add_edges_from(map(lambda a, b, data: (f(a), f(b), data), *zip(*graph.edges_iter(data=True))))
return digraph
def merge_nodes(graph):
2015-03-24 16:49:37 +01:00
def merge_edges(data):
tq = min(map(lambda d: d['tq'], data))
vpn = all(map(lambda d: d.get('vpn', False), data))
return dict(tq=tq, vpn=vpn)
multigraph = to_multigraph(graph)
digraph = nx.DiGraph()
digraph.add_nodes_from(multigraph.nodes_iter(data=True))
edges = chain.from_iterable([[(e, d, merge_edges(multigraph[e][d].values()))
for d in multigraph[e]] for e in multigraph])
digraph.add_edges_from(edges)
2015-03-24 16:49:37 +01:00
return digraph
def to_undirected(graph):
2015-03-24 16:49:37 +01:00
multigraph = nx.MultiGraph()
multigraph.add_nodes_from(graph.nodes_iter(data=True))
multigraph.add_edges_from(graph.edges_iter(data=True))
2015-03-24 16:49:37 +01:00
def merge_edges(data):
tq = max(map(lambda d: d['tq'], data))
vpn = all(map(lambda d: d.get('vpn', False), data))
return dict(tq=tq, vpn=vpn, bidirect=len(data) == 2)
2015-03-24 16:49:37 +01:00
graph = nx.Graph()
graph.add_nodes_from(multigraph.nodes_iter(data=True))
edges = chain.from_iterable([[(e, d, merge_edges(multigraph[e][d].values()))
for d in multigraph[e]] for e in multigraph])
graph.add_edges_from(edges)
2015-03-24 16:49:37 +01:00
return graph