from graphviz import Digraph from collections import namedtuple class NetworkGraph: ''' Representation of the network connections. This class contains the entities in the network e.g. hosts or switches. And the connections between them. ''' Vertex = namedtuple('Vertexes', ['hosts', 'switches']) _edges = [] def _sanitize_edge_connection(self, edge): ''' Update '_to' and '_form' field of a edge. :param edge: One edge connection. :type edge: dict :returns: Updated edge with _to and _from key. :rtype: dict ''' if edge['to_guid'].startswith('S'): to_collection = 'switches/' elif edge['to_guid'].startswith('H'): to_collection = 'hosts/' if edge['from_guid'].startswith('S'): from_collection = 'switches/' elif edge['from_guid'].startswith('H'): from_collection = 'hosts/' edge.update({ '_to': to_collection + edge['to_guid'], '_from': from_collection + edge['from_guid'] }) return edge def _sanitize_vertexes(self, vertex): ''' Update '_key' field of vertex to appropriate guid. :param vertex: Vertex :type vertex: dict :returns: An updated dict, '_key' field with 'guid' value. :rtype: dict ''' vertex.update({'_key': vertex['guid']}) return vertex def __init__(self, hsts=None, switches=None, connections=None): self._vertexes = self.Vertex(hosts=[self._sanitize_vertexes(h) for h in hsts], switches=[self._sanitize_vertexes(s) for s in switches]) self._edges = [self._sanitize_edge_connection(c) for c in connections] @property def vertexes(self): ''' Returns a concatenated list of all vertexes. :returns: List of vertexes, contains of hosts and switches. :rtype: List[dict] ''' return self._vertexes.hosts + self._vertexes.switches @property def switches(self): ''' Returns a list of all 'switch' vertexes. :returns: List of all switches. :rtype: List[dict] ''' return self._vertexes.switches @property def hosts(self): ''' Returns a list of all 'host' vertexes. :returns: List of all hosts. :rtype: List[dict] ''' return self._vertexes.hosts @property def edges(self): ''' Return a list of all 'connection' edges. :returns: List of all connections. :rtype: List[dict] ''' return self._edges def to_graph(self, graphargs): ''' Draw a dot graph of the network graph. :params graphargs: Arguments to graphviz.Digraph. :type graphargs: dict ''' graph = Digraph(**graphargs) for v in self._vertexes: graph.node(v['guid'], v['description']) for c in self._edges: graph.edge(c['from_guid'], c['to_guid']) graph.render()