Refactor D2 graph generation

Implementing ideas from #2535, add code that creates generic graph
data structure for topology and BGP graphs, and use that data structure
to generate D2 graphs with generic primitives (d2_cluster, d2_node, and
d2_edge)

Author: ipspace

netsim/outputs/_graph.py

@@ -0,0 +1,148 @@
+#
+# Create generic graph from lab topology
+#
+import typing
+from box import Box
+
+from ..data import get_box,get_empty_box
+from ..data.types import must_be_list
+from ..utils import log
+
+def build_nodes(topology: Box) -> Box:
+  maps = Box({},default_box=True,box_dots=True)
+  for name,n in topology.nodes.items():
+    maps.nodes[name] = n
+
+  if 'bgp' in topology.get('module',[]):
+    for name,n in topology.nodes.items():
+      bgp_as = n.get('bgp.as',None)
+      if bgp_as:
+        bgp_as = f'AS_{bgp_as}'
+        maps.bgp[bgp_as].nodes[n.name] = n
+
+  if 'bgp' in topology and 'as_list' in topology.bgp:
+    for (asn,asdata) in topology.bgp.as_list.items():
+      if 'name' in asdata and asn in maps.bgp:
+        maps.bgp[asn].name = asdata.name
+
+  return maps
+
+'''
+add_groups -- use topology groups as graph clustering mechanism
+'''
+def add_groups(maps: Box, graph_groups: list, topology: Box) -> None:
+  if not 'groups' in topology:
+    return
+  placed_hosts = []
+
+  for g_name,g_data in topology.groups.items():
+    if g_name not in graph_groups:
+      continue
+    for node in g_data.members:
+      if node in placed_hosts:
+        log.error(
+          f'Cannot create overlapping graph clusters: node {node} is in two groups',
+          log.IncorrectValue,
+          'graph')
+        continue
+
+      maps.clusters[g_name].nodes[node] = topology.nodes[node]
+      placed_hosts.append(node)
+
+def graph_clusters(graph: Box, topology: Box, settings: Box) -> None:
+  if 'groups' in settings:
+    must_be_list(
+      parent=settings,
+      key='groups',path='defaults.outputs.graph',
+      true_value=list(topology.get('groups',{}).keys()),
+      create_empty=True,
+      module='graph')
+    add_groups(graph,settings.groups,topology)
+  elif 'bgp' in graph and settings.as_clusters:
+    graph.clusters = graph.bgp
+
+def append_edge(graph: Box, if_a: Box, if_b: Box, g_type: str) -> None:
+  nodes = []
+  e_attr = get_empty_box()
+  for intf in (if_a, if_b):
+    addr = intf.ipv4 or intf.ipv6
+    if addr and 'prefix' not in intf:
+      addr = addr.split('/')[0]
+    intf_attr = intf.get('graph',{}) + intf.get(g_type,{})
+    e_attr += intf_attr
+    e_data = get_box({
+               'node': intf.node,
+               'attr': intf_attr,
+               'label': addr})
+    for kw in ['type']:
+      if kw in intf:
+        e_data[kw] = intf[kw]
+
+    nodes.append(e_data)
+
+  graph.edges.append({'nodes': nodes, 'attr': e_attr})
+
+def topo_edges(graph: Box, topology: Box, settings: Box,g_type: str) -> None:
+  graph.edges = []
+  for link in sorted(topology.links,key=lambda x: x.get(f'{g_type}.linkorder',100)):
+    l_attr = link.get(g_type,None)
+    if l_attr:
+      for intf in link.interfaces:
+        intf[g_type] = l_attr + intf[g_type]
+
+    if len(link.interfaces) == 2:
+      append_edge(graph,link.interfaces[0],link.interfaces[1],g_type)
+    else:
+      link.node = link.name or f'{link.type}_{link.linkindex}'
+      graph.nodes[link.node] = link
+      for af in ['ipv4','ipv6']:
+        if af in link.prefix:
+          link[af] = link.prefix[af]
+
+      for intf in link.interfaces:
+        e_data = [ intf, link ]
+        if intf.get(f'{g_type}.linkorder',100) > link.get(f'{g_type}.linkorder',101):
+          e_data = [ link, intf ]
+        append_edge(graph,e_data[0],e_data[1],g_type)
+
+def topology_graph(topology: Box, settings: Box,g_type: str) -> Box:
+  graph = build_nodes(topology)
+  graph_clusters(graph,topology,settings)
+  topo_edges(graph,topology,settings,g_type)
+  return graph
+
+def bgp_sessions(graph: Box, topology: Box, settings: Box, g_type: str, rr_sessions: bool) -> None:
+  for n_name,n_data in topology.nodes.items():
+    if 'bgp' not in n_data:
+      continue
+    for neighbor in n_data.bgp.get('neighbors',[]):
+      if neighbor.name < n_name:
+        continue
+
+      e_1 = get_box({ 'node': n_name, 'type': neighbor.type })
+      e_2 = get_box({ 'node': neighbor.name, 'type': neighbor.type })
+      dir = '<->'
+      if 'ibgp' in neighbor.type:
+        if n_data.bgp.get('rr',False) and not neighbor.get('rr',False):
+          dir = '->'
+        elif neighbor.get('rr',False) and not n_data.get('rr',False):
+          dir = '->'
+          e_1, e_2 = e_2, e_1
+
+      if rr_sessions:
+        e_1.graph.dir = dir
+      append_edge(graph,e_1,e_2,g_type)
+
+def bgp_graph(topology: Box, settings: Box, g_type: str, rr_sessions: bool) -> typing.Optional[Box]:
+  if 'bgp' not in topology.module:
+    log.error(
+      'Cannot build a BGP graph from a topology that does not use BGP',
+      category=log.IncorrectType,
+      module='graph')
+    return None
+
+  graph = build_nodes(topology)
+  graph.clusters = graph.bgp
+  graph.edges = []
+  bgp_sessions(graph,topology,settings,g_type,rr_sessions)
+  return graph

netsim/outputs/d2.py

@@ -12,6 +12,7 @@
 from . import _TopologyOutput
 from ..utils import files as _files
 from ..utils import log
+from ._graph import topology_graph,bgp_graph
 
 '''
 Copy default settings into a D2 map converting Python dictionaries into
@@ -47,13 +48,16 @@ def d2_node_attr(f : typing.TextIO, n: Box, settings: Box, indent: str = '') ->
 Add D2 styling information from d2.* link/node attributes
 '''
 STYLE_MAP: Box
+IGNORE_KW: list = ['dir', 'type', 'name']
 
 def d2_style(f : typing.TextIO, obj: Box, indent: str) -> None:
   if 'd2' not in obj:
     return
-  d2_data = { STYLE_MAP[k]:v for k,v in obj.d2.items() if k in STYLE_MAP }
-  if d2_data:
-    dump_d2_dict(f,{ 'style': d2_data },indent)
+  d2_style = { STYLE_MAP[k]:v for k,v in obj.d2.items() if k in STYLE_MAP }
+  d2_extra = get_box({ k:v for k,v in obj.d2.items() if k not in STYLE_MAP and k not in IGNORE_KW })
+
+  if d2_style or d2_extra:
+    dump_d2_dict(f,d2_extra + { 'style': d2_style },indent)
 
 '''
 Create a node in D2 graph and add a label and styling attributes to it
@@ -108,161 +112,90 @@ def edge_label(f : typing.TextIO, direction: str, data: Box, subnet: bool = True
 Create a P2P connection between two nodes
 '''
 def edge_p2p(f : typing.TextIO, l: Box, labels: typing.Optional[bool] = False) -> None:
-  f.write(f"{l.interfaces[0].node} -- {l.interfaces[1].node} {{\n")
+  e_direction = ('source','target')
+  dir = l.interfaces[0].get('attr.dir','--')
+  f.write(f"{l.interfaces[0].d2.name} {dir} {l.interfaces[1].d2.name} {{\n")
+  print(l)
   d2_style(f,l,'  ')
   if labels:
-    edge_label(f,'source',l.interfaces[0],True)
-    edge_label(f,'target',l.interfaces[1],True)
+    for e_idx,intf in enumerate(l.interfaces):
+      if 'prefix' not in intf:
+        edge_label(f,e_direction[e_idx],intf,True)
   f.write("}\n")
 
 '''
-Create a connection between a node and a LAN segment
+Create a group container (or ASN container)
 '''
-def edge_node_net(f : typing.TextIO, link: Box, ifdata: Box, labels: typing.Optional[bool] = False) -> None:
-  f.write(f"{ifdata.node} -> {link.bridge} {{\n")
-  d2_style(f,link,'  ')
-  if labels:
-    edge_label(f,'source',ifdata,False)
-  f.write("}\n")
-
-'''
-Create an ASN or group container
-'''
-def as_start(f : typing.TextIO, asn: str, label: typing.Optional[str], settings: Box) -> None:
+def d2_cluster_start(f : typing.TextIO, asn: str, label: typing.Optional[str], settings: Box) -> None:
   f.write(f'{asn} {{\n')
   copy_d2_attr(f,'container',settings,'-  ')
   asn = asn.replace('_',' ')
   f.write('  label: '+ (f'{label} ({asn})' if label else asn)+'\n')
 
 '''
-Create a topology graph as a set of containers (groups or ASNs)
+Create graph containers
 '''
-def graph_clusters(f : typing.TextIO, clusters: Box, settings: Box, nodes: Box) -> None:
-  for asn in clusters.keys():
-    as_start(f,asn,clusters[asn].get('name',None),settings)
-    for n in clusters[asn].nodes.values():
+def d2_clusters(f: typing.TextIO, graph: Box, topology: Box, settings: Box) -> None:
+  for c_name,c_data in graph.clusters.items():
+    d2_cluster_start(f,c_name,c_data.get('name',None),settings)
+    for n in c_data.nodes.values():
       node_with_label(f,n,settings,'  ')                    # Create a node within a cluster
-      nodes[n.name].d2.name = f'{asn}.{n.name}'             # And change the D2 node name that will be used in connections
+      topology.nodes[n.name].d2.name = f'{c_name}.{n.name}' # And change the D2 node name that will be used in connections
     f.write('}\n')
 
-def build_maps(topology: Box) -> Box:
-  maps = Box({},default_box=True,box_dots=True)
-  for name,n in topology.nodes.items():
-    maps.nodes[name] = n
-
-  if 'bgp' in topology.get('module',[]):
-    for name,n in topology.nodes.items():
-      bgp_as = n.get('bgp.as',None)
-      if bgp_as:
-        bgp_as = f'AS_{bgp_as}'
-        maps.bgp[bgp_as].nodes[n.name] = n
-
-  if 'bgp' in topology and 'as_list' in topology.bgp:
-    for (asn,asdata) in topology.bgp.as_list.items():
-      if 'name' in asdata and asn in maps.bgp:
-        maps.bgp[asn].name = asdata.name
-
-  return maps
-
-'''
-add_groups -- use topology groups as graph clustering mechanism
-'''
-def add_groups(maps: Box, groups: list, topology: Box) -> None:
-  if not 'groups' in topology:
-    return
-
-  placed_hosts = []
-  for g,v in topology.groups.items():
-    if g in groups:
-      for n in v.members:
-        if n in placed_hosts:
-          log.error(
-            f'Cannot create overlapping graph clusters: node {n} is in two groups',
-            log.IncorrectValue,
-            'graph')
-          continue
-        else:
-          maps.groups[g].nodes[n] = topology.nodes[n]
-          placed_hosts.append(n)
+def d2_nodes(f: typing.TextIO, graph: Box, topology: Box, settings: Box) -> None:
+  for n_name,n_data in graph.nodes.items():
+    if 'd2.name' not in n_data:
+      n_data.d2.name = n_data.name
+      if 'prefix' in n_data:
+        network_with_label(f,n_data,settings)
+      else:
+        node_with_label(f,n_data,settings)
+
+def d2_links(f: typing.TextIO, graph: Box, topology: Box, settings: Box) -> None:
+  for edge in graph.edges:
+    fake_link = get_box({ 'interfaces': edge.nodes, 'd2': edge.attr })
+    for intf in edge.nodes:
+      intf.d2.name = intf.node
+      if intf.node in topology.nodes:
+        intf.d2.name = topology.nodes[intf.node].d2.name
+    edge_p2p(f,fake_link,settings.interface_labels)
 
 def graph_topology(topology: Box, fname: str, settings: Box,g_format: typing.Optional[list]) -> bool:
+  graph = topology_graph(topology,settings,'d2')
   f = _files.open_output_file(fname)
-  maps = build_maps(topology)
-
-  if 'groups' in settings:
-    must_be_list(
-      parent=settings,
-      key='groups',path='defaults.outputs.graph',
-      true_value=list(topology.get('groups',{}).keys()),
-      create_empty=True,
-      module='graph')
-    add_groups(maps,settings.groups,topology)
-    graph_clusters(f,maps.groups,settings,topology.nodes)
-  elif 'bgp' in maps and settings.as_clusters:
-    graph_clusters(f,maps.bgp,settings,topology.nodes)
-  else:
-    for name,n in topology.nodes.items():
-      node_with_label(f,n,settings)
-
-  for l in sorted(topology.links,key=lambda x: x.get('d2.linkorder',100)):
-    for intf in l.interfaces:
-      intf._topo_node = intf.node
-      intf.node = topology.nodes[intf.node].d2.name
-    if l.type == "p2p":
-      edge_p2p(f,l,settings.interface_labels)
-    else:
-      l.bridge = l.name or f'{l.type}_{l.linkindex}'
-      network_with_label(f,l,settings)
-      for ifdata in l.interfaces:
-        if ifdata._topo_node in maps.nodes:
-          edge_node_net(f,l,ifdata,settings.interface_labels)
 
-  f.close()
-  return True
+  d2_clusters(f,graph,topology,settings)
+  d2_nodes(f,graph,topology,settings)
+  d2_links(f,graph,topology,settings)
 
-'''
-Create a BGP session as a connection between two nodes
+  if fname != '-':
+    f.close()
+  return True
 
-* Select the connection (arrow) type based on whether a connection is a RR-to-client session
-* Copy IBGP or EBGP attributes to the connection
+def graph_bgp(topology: Box, fname: str, settings: Box, g_format: typing.Optional[list]) -> bool:
+  rr_session = settings.get('rr_sessions',False)
+  if g_format is not None and len(g_format) > 1:
+    rr_session = g_format[1] == 'rr'
 
-Please note that we call this function once for every pair of nodes, so it has to deal with
-RR being the first (node) or the second (session) connection endpoint.
-'''
-def bgp_session(f : typing.TextIO, node: Box, session: Box, settings: Box, rr_session: bool, nodes: Box) -> None:
-  arrow_dir = '<->'
-  if rr_session:
-    if session.type == 'ibgp':
-      if 'rr' in node.bgp and node.bgp.rr and not 'rr' in session:
-        arrow_dir = '->'
-      if not 'rr' in node.bgp and 'rr' in session:
-        arrow_dir = '<-'
-
-  f.write(f'{node.d2.name} {arrow_dir} {nodes[session.name].d2.name} {{\n')
-  copy_d2_attr(f,session.type,settings,'  ')
-  f.write('}\n')
-
-def graph_bgp(topology: Box, fname: str, settings: Box,g_format: typing.Optional[list]) -> bool:
-  if not 'bgp' in topology.get('module',{}):
-    log.error('BGP graph format can only be used to draw topologies using BGP',module='d2')
+  graph = bgp_graph(topology,settings,'d2',rr_sessions=rr_session)
+  if graph is None:
     return False
 
   f = _files.open_output_file(fname)
 
-  rr_session = settings.get('rr_sessions',False)
-  if g_format is not None and len(g_format) > 1:
-    rr_session = g_format[1] == 'rr'
+  d2_clusters(f,graph,topology,settings)
+  d2_nodes(f,graph,topology,settings)
 
-  maps = build_maps(topology)
-  graph_clusters(f,maps.bgp,settings,topology.nodes)
+  for edge in graph.edges:
+    print(edge)
+    if edge.nodes[0].type in settings:
+      edge.attr = settings[edge.nodes[0].type] + edge.attr
 
-  for name,n in topology.nodes.items():
-    if 'bgp' in n:
-      for neighbor in n.bgp.get('neighbors',[]):
-        if neighbor.name > n.name:
-          bgp_session(f,n,neighbor,settings,rr_session,topology.nodes)
+  d2_links(f,graph,topology,settings)
 
-  f.close()
+  if fname != '-':
+    f.close()
   return True
 
 graph_dispatch = {