python/samba/graph.py

   1 # -*- coding: utf-8 -*-
   2 # Graph topology utilities and dot file generation
   3 #
   4 # Copyright (C) Andrew Bartlett 2018.
   5 #
   6 # Written by Douglas Bagnall <douglas.bagnall@catalyst.net.nz>
   7 #
   8 # This program is free software; you can redistribute it and/or modify
   9 # it under the terms of the GNU General Public License as published by
  10 # the Free Software Foundation; either version 3 of the License, or
  11 # (at your option) any later version.
  12 #
  13 # This program is distributed in the hope that it will be useful,
  14 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  15 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  16 # GNU General Public License for more details.
  17 #
  18 # You should have received a copy of the GNU General Public License
  19 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  20
  21 from __future__ import print_function
  22 from samba import colour
  23 import sys
  24
  25 FONT_SIZE = 10
  26
  27
  28 def reformat_graph_label(s):
  29     """Break DNs over multiple lines, for better shaped and arguably more
  30     readable nodes. We try to split after commas, and if necessary
  31     after hyphens or failing that in arbitrary places."""
  32     if len(s) < 12:
  33         return s
  34
  35     s = s.replace(',', ',\n')
  36     pieces = []
  37     for p in s.split('\n'):
  38         while len(p) > 20:
  39             if '-' in p[2:20]:
  40                 q, p = p.split('-', 1)
  41             else:
  42                 n = len(p) / 12
  43                 b = len(p) / n
  44                 q, p = p[:b], p[b:]
  45             pieces.append(q + '-')
  46         if p:
  47             pieces.append(p)
  48
  49     return '\\n'.join(pieces)
  50
  51
  52 def quote_graph_label(s, reformat=False):
  53     """Escape a string as graphvis requires."""
  54     # escaping inside quotes is simple in dot, because only " is escaped.
  55     # there is no need to count backslashes in sequences like \\\\"
  56     s = s.replace('"', '\"')
  57     if reformat:
  58         s = reformat_graph_label(s)
  59     return "%s" % s
  60
  61
  62 def shorten_vertex_names(edges, vertices, suffix=',...', aggressive=False):
  63     """Replace the common suffix (in practice, the base DN) of a number of
  64     vertices with a short string (default ",..."). If this seems
  65     pointless because the replaced string is very short or the results
  66     seem strange, the original vertices are retained.
  67
  68     :param edges: a sequence of vertex pairs to shorten
  69     :param vertices: a sequence of vertices to shorten
  70     :param suffix: the replacement string [",..."]
  71
  72     :return: tuple of (edges, vertices, replacement)
  73
  74     If no change is made, the returned edges and vertices will be the
  75     original lists  and replacement will be None.
  76
  77     If a change is made, replacement will be a tuple (new, original)
  78     indicating the new suffix that replaces the old.
  79     """
  80     vlist = list(set(x[0] for x in edges) |
  81                  set(x[1] for x in edges) |
  82                  set(vertices))
  83
  84     if len(vlist) < 2:
  85         return edges, vertices, None
  86
  87     # walk backwards along all the strings until we meet a character
  88     # that is not shared by all.
  89     i = -1
  90     try:
  91         while True:
  92             c = set(x[i] for x in vlist)
  93             if len(c) > 1:
  94                 break
  95             i -= 1
  96     except IndexError:
  97         # We have indexed beyond the start of a string, which should
  98         # only happen if one node is a strict suffix of all others.
  99         return edges, vertices, None
 100
 101     # add one to get to the last unanimous character.
 102     i += 1
 103
 104     # now, we actually really want to split on a comma. So we walk
 105     # back to a comma.
 106     x = vlist[0]
 107     while i < len(x) and x[i] != ',':
 108         i += 1
 109
 110     if i >= -len(suffix):
 111         # there is nothing to gain here
 112         return edges, vertices, None
 113
 114     edges2 = []
 115     vertices2 = []
 116
 117     for a, b in edges:
 118         edges2.append((a[:i] + suffix, b[:i] + suffix))
 119     for a in vertices:
 120         vertices2.append(a[:i] + suffix)
 121
 122     replacements = [(suffix, a[i:])]
 123
 124     if aggressive:
 125         # Remove known common annoying strings
 126         map = dict((v, v) for v in vertices2)
 127         for v in vertices2:
 128             if ',CN=Servers,' not in v:
 129                 break
 130         else:
 131             map = dict((k, v.replace(',CN=Servers,', ',**,'))
 132                        for k, v in map.iteritems())
 133             replacements.append(('**', 'CN=Servers'))
 134
 135         for v in vertices2:
 136             if not v.startswith('CN=NTDS Settings,'):
 137                 break
 138         else:
 139             map = dict((k, v.replace('CN=NTDS Settings,', '*,'))
 140                        for k, v in map.iteritems())
 141             replacements.append(('*', 'CN=NTDS Settings'))
 142
 143         edges2 = [(map.get(a, a), map.get(b, b)) for a, b in edges2]
 144         vertices2 = [map.get(a, a) for a in vertices2]
 145
 146     return edges2, vertices2, replacements
 147
 148
 149 def compile_graph_key(key_items, nodes_above=[], elisions=None,
 150                       prefix='key_', width=2):
 151     """Generate a dot file snippet that acts as a legend for a graph.
 152
 153     :param key_items: sequence of items (is_vertex, style, label)
 154     :param nodes_above: list of vertices (pushes key into right position)
 155     :param elision: tuple (short, full) indicating suffix replacement
 156     :param prefix: string used to generate key node names ["key_"]
 157     :param width: default width of node lines
 158
 159     Each item in key_items is a tuple of (is_vertex, style, label).
 160     is_vertex is a boolean indicating whether the item is a vertex
 161     (True) or edge (False). Style is a dot style string for the edge
 162     or vertex. label is the text associated with the key item.
 163     """
 164     edge_lines = []
 165     edge_names = []
 166     vertex_lines = []
 167     vertex_names = []
 168     order_lines = []
 169     for i, item in enumerate(key_items):
 170         is_vertex, style, label = item
 171         tag = '%s%d_' % (prefix, i)
 172         label = quote_graph_label(label)
 173         name = '%s_label' % tag
 174
 175         if is_vertex:
 176             order_lines.append(name)
 177             vertex_names.append(name)
 178             vertex_lines.append('%s[label="%s"; %s]' %
 179                                 (name, label, style))
 180         else:
 181             edge_names.append(name)
 182             e1 = '%se1' % tag
 183             e2 = '%se2' % tag
 184             order_lines.append(name)
 185             edge_lines.append('subgraph cluster_%s {' % tag)
 186             edge_lines.append('%s[label=src; color="#000000"; group="%s_g"]' %
 187                               (e1, tag))
 188             edge_lines.append('%s[label=dest; color="#000000"; group="%s_g"]' %
 189                               (e2, tag))
 190             edge_lines.append('%s -> %s [constraint = false; %s]' % (e1, e2,
 191                                                                      style))
 192             edge_lines.append(('%s[shape=plaintext; style=solid; width=%f; '
 193                                'label="%s\\r"]') %
 194                               (name, width, label))
 195             edge_lines.append('}')
 196
 197     elision_str = ''
 198     if elisions:
 199         for i, elision in enumerate(reversed(elisions)):
 200             order_lines.append('elision%d' % i)
 201             short, long = elision
 202             if short[0] == ',' and long[0] == ',':
 203                 short = short[1:]
 204                 long = long[1:]
 205             elision_str += ('\nelision%d[shape=plaintext; style=solid; '
 206                             'label="\“%s”  means  “%s”\\r"]\n'
 207                             % ((i, short, long)))
 208
 209     above_lines = []
 210     if order_lines:
 211         for n in nodes_above:
 212             above_lines.append('"%s" -> %s [style=invis]' %
 213                                (n, order_lines[0]))
 214
 215     s = ('subgraph cluster_key {\n'
 216          'label="Key";\n'
 217          'subgraph cluster_key_nodes {\n'
 218          'label="";\n'
 219          'color = "invis";\n'
 220          '%s\n'
 221          '}\n'
 222          'subgraph cluster_key_edges {\n'
 223          'label="";\n'
 224          'color = "invis";\n'
 225          '%s\n'
 226          '{%s}\n'
 227          '}\n'
 228          '%s\n'
 229          '}\n'
 230          '%s\n'
 231          '%s [style=invis; weight=9]'
 232          '\n'
 233          % (';\n'.join(vertex_lines),
 234             '\n'.join(edge_lines),
 235             ' '.join(edge_names),
 236             elision_str,
 237             ';\n'.join(above_lines),
 238             ' -> '.join(order_lines),
 239          ))
 240
 241     return s
 242
 243
 244 def dot_graph(vertices, edges,
 245               directed=False,
 246               title=None,
 247               reformat_labels=True,
 248               vertex_colors=None,
 249               edge_colors=None,
 250               edge_labels=None,
 251               vertex_styles=None,
 252               edge_styles=None,
 253               graph_name=None,
 254               shorten_names=False,
 255               key_items=None,
 256               vertex_clusters=None):
 257     """Generate a Graphviz representation of a list of vertices and edges.
 258
 259     :param vertices: list of vertex names (optional).
 260     :param edges:    list of (vertex, vertex) pairs
 261     :param directed: bool: whether the graph is directed
 262     :param title: optional title for the graph
 263     :param reformat_labels: whether to wrap long vertex labels
 264     :param vertex_colors: if not None, a sequence of colours for the vertices
 265     :param edge_colors: if not None, colours for the edges
 266     :param edge_labels: if not None, labels for the edges
 267     :param vertex_styles: if not None, DOT style strings for vertices
 268     :param edge_styles: if not None, DOT style strings for edges
 269     :param graph_name: if not None, name of graph
 270     :param shorten_names: if True, remove common DN suffixes
 271     :param key: (is_vertex, style, description) tuples
 272     :param vertex_clusters: list of subgraph cluster names
 273
 274     Colour, style, and label lists must be the same length as the
 275     corresponding list of edges or vertices (or None).
 276
 277     Colours can be HTML RGB strings ("#FF0000") or common names
 278     ("red"), or some other formats you don't want to think about.
 279
 280     If `vertices` is None, only the vertices mentioned in the edges
 281     are shown, and their appearance can be modified using the
 282     vertex_colors and vertex_styles arguments. Vertices appearing in
 283     the edges but not in the `vertices` list will be shown but their
 284     styles can not be modified.
 285     """
 286     out = []
 287     write = out.append
 288
 289     if vertices is None:
 290         vertices = set(x[0] for x in edges) | set(x[1] for x in edges)
 291
 292     if shorten_names:
 293         edges, vertices, elisions = shorten_vertex_names(edges, vertices)
 294     else:
 295         elisions = None
 296
 297     if graph_name is None:
 298         graph_name = 'A_samba_tool_production'
 299
 300     if directed:
 301         graph_type = 'digraph'
 302         connector = '->'
 303     else:
 304         graph_type = 'graph'
 305         connector = '--'
 306
 307     write('/* generated by samba */')
 308     write('%s %s {' % (graph_type, graph_name))
 309     if title is not None:
 310         write('label="%s";' % (title,))
 311     write('fontsize=%s;\n' % (FONT_SIZE))
 312     write('node[fontname=Helvetica; fontsize=%s];\n' % (FONT_SIZE))
 313
 314     prev_cluster = None
 315     cluster_n = 0
 316     quoted_vertices = []
 317     for i, v in enumerate(vertices):
 318         v = quote_graph_label(v, reformat_labels)
 319         quoted_vertices.append(v)
 320         attrs = []
 321         if vertex_clusters and vertex_clusters[i]:
 322             cluster = vertex_clusters[i]
 323             if cluster != prev_cluster:
 324                 if prev_cluster is not None:
 325                     write("}")
 326                 prev_cluster = cluster
 327                 n = quote_graph_label(cluster)
 328                 if cluster:
 329                     write('subgraph cluster_%d {' % cluster_n)
 330                     cluster_n += 1
 331                     write('style = "rounded,dotted";')
 332                     write('node [style="filled"; fillcolor=white];')
 333                     write('label = "%s";' % n)
 334
 335         if vertex_styles and vertex_styles[i]:
 336             attrs.append(vertex_styles[i])
 337         if vertex_colors and vertex_colors[i]:
 338             attrs.append('color="%s"' % quote_graph_label(vertex_colors[i]))
 339         if attrs:
 340             write('"%s" [%s];' % (v, ', '.join(attrs)))
 341         else:
 342             write('"%s";' % (v,))
 343
 344     if prev_cluster:
 345         write("}")
 346
 347     for i, edge in enumerate(edges):
 348         a, b = edge
 349         if a is None:
 350             a = "Missing source value"
 351         if b is None:
 352             b = "Missing destination value"
 353
 354         a = quote_graph_label(a, reformat_labels)
 355         b = quote_graph_label(b, reformat_labels)
 356
 357         attrs = []
 358         if edge_labels:
 359             label = quote_graph_label(edge_labels[i])
 360             attrs.append('label="%s"' % label)
 361         if edge_colors:
 362             attrs.append('color="%s"' % quote_graph_label(edge_colors[i]))
 363         if edge_styles:
 364             attrs.append(edge_styles[i])  # no quoting
 365         if attrs:
 366             write('"%s" %s "%s" [%s];' % (a, connector, b, ', '.join(attrs)))
 367         else:
 368             write('"%s" %s "%s";' % (a, connector, b))
 369
 370     if key_items:
 371         key = compile_graph_key(key_items, nodes_above=quoted_vertices,
 372                                 elisions=elisions)
 373         write(key)
 374
 375     write('}\n')
 376     return '\n'.join(out)
 377
 378
 379 COLOUR_SETS = {
 380     'ansi': {
 381         'alternate rows': (colour.DARK_WHITE, colour.BLACK),
 382         'disconnected': colour.RED,
 383         'connected': colour.GREEN,
 384         'transitive': colour.DARK_YELLOW,
 385         'header': colour.UNDERLINE,
 386         'reset': colour.C_NORMAL,
 387     },
 388     'ansi-heatmap': {
 389         'alternate rows': (colour.DARK_WHITE, colour.BLACK),
 390         'disconnected': colour.REV_RED,
 391         'connected': colour.REV_GREEN,
 392         'transitive': colour.REV_DARK_YELLOW,
 393         'header': colour.UNDERLINE,
 394         'reset': colour.C_NORMAL,
 395     },
 396     'xterm-256color': {
 397         'alternate rows': (colour.xterm_256_colour(39),
 398                            colour.xterm_256_colour(45)),
 399         #'alternate rows': (colour.xterm_256_colour(246),
 400         #                   colour.xterm_256_colour(247)),
 401         'disconnected': colour.xterm_256_colour(124, bg=True),
 402         'connected': colour.xterm_256_colour(112),
 403         'transitive': colour.xterm_256_colour(214),
 404         'transitive scale': (colour.xterm_256_colour(190),
 405                              colour.xterm_256_colour(226),
 406                              colour.xterm_256_colour(220),
 407                              colour.xterm_256_colour(214),
 408                              colour.xterm_256_colour(208),
 409         ),
 410         'header': colour.UNDERLINE,
 411         'reset': colour.C_NORMAL,
 412     },
 413     'xterm-256color-heatmap': {
 414         'alternate rows': (colour.xterm_256_colour(171),
 415                            colour.xterm_256_colour(207)),
 416         #'alternate rows': (colour.xterm_256_colour(246),
 417         #                    colour.xterm_256_colour(247)),
 418         'disconnected': colour.xterm_256_colour(124, bg=True),
 419         'connected': colour.xterm_256_colour(112, bg=True),
 420         'transitive': colour.xterm_256_colour(214, bg=True),
 421         'transitive scale': (colour.xterm_256_colour(190, bg=True),
 422                              colour.xterm_256_colour(226, bg=True),
 423                              colour.xterm_256_colour(220, bg=True),
 424                              colour.xterm_256_colour(214, bg=True),
 425                              colour.xterm_256_colour(208, bg=True),
 426         ),
 427         'header': colour.UNDERLINE,
 428         'reset': colour.C_NORMAL,
 429     },
 430     None: {
 431         'alternate rows': ('',),
 432         'disconnected': '',
 433         'connected': '',
 434         'transitive': '',
 435         'header': '',
 436         'reset': '',
 437     }
 438 }
 439
 440
 441 def find_transitive_distance(vertices, edges):
 442     all_vertices = (set(vertices) |
 443                     set(e[0] for e in edges) |
 444                     set(e[1] for e in edges))
 445
 446     if all_vertices != set(vertices):
 447         print("there are unknown vertices: %s" %
 448               (all_vertices - set(vertices)),
 449               file=sys.stderr)
 450
 451     # with n vertices, we are always less than n hops away from
 452     # anywhere else.
 453     inf = len(all_vertices)
 454     distances = {}
 455     for v in all_vertices:
 456         distances[v] = {v: 0}
 457
 458     for src, dest in edges:
 459         distances[src][dest] = distances[src].get(dest, 1)
 460
 461     # This algorithm (and implementation) seems very suboptimal.
 462     # potentially O(n^4), though n is smallish.
 463     for i in range(inf):
 464         changed = False
 465         new_distances = {}
 466         for v, d in distances.iteritems():
 467             new_d = d.copy()
 468             new_distances[v] = new_d
 469             for dest, cost in d.iteritems():
 470                 for leaf, cost2 in distances[dest].iteritems():
 471                     new_cost = cost + cost2
 472                     old_cost = d.get(leaf, inf)
 473                     if new_cost < old_cost:
 474                         new_d[leaf] = new_cost
 475                         changed = True
 476
 477         distances = new_distances
 478         if not changed:
 479             break
 480
 481     # filter out unwanted vertices and infinite links
 482     answer = {}
 483     for v in vertices:
 484         answer[v] = {}
 485         for v2 in vertices:
 486             a = distances[v].get(v2, inf)
 487             if a < inf:
 488                 answer[v][v2] = a
 489
 490     return answer
 491
 492
 493 def get_transitive_colourer(colours, n_vertices):
 494     if 'transitive scale' in colours:
 495         scale = colours['transitive scale']
 496         m = len(scale)
 497         n = 1 + int(n_vertices ** 0.5)
 498
 499         def f(link):
 500             return scale[min(link * m / n, m - 1)]
 501
 502     else:
 503         def f(link):
 504             return colours['transitive']
 505
 506     return f
 507
 508
 509 def distance_matrix(vertices, edges,
 510                     utf8=False,
 511                     colour=None,
 512                     shorten_names=False,
 513                     generate_key=False):
 514     lines = []
 515     write = lines.append
 516
 517     if utf8:
 518         vertical = '│'
 519         horizontal = '─'
 520         corner = '╭'
 521         #diagonal = '╲'
 522         diagonal = '·'
 523         #missing = '🕱'
 524         missing = '-'
 525     else:
 526         vertical, horizontal, corner, diagonal, missing = '|-,0-'
 527
 528     colours = COLOUR_SETS[colour]
 529
 530     if vertices is None:
 531         vertices = sorted(set(x[0] for x in edges) | set(x[1] for x in edges))
 532
 533     if shorten_names:
 534         edges, vertices, replacements = shorten_vertex_names(edges,
 535                                                              vertices,
 536                                                              '+',
 537                                                              aggressive=True)
 538
 539     vlen = max(6, max(len(v) for v in vertices))
 540
 541     # first, the key for the columns
 542     colour_cycle = colours.get('alternate rows', ('',))
 543     c_header = colours.get('header', '')
 544     c_disconn = colours.get('disconnected', '')
 545     c_conn = colours.get('connected', '')
 546     c_reset = colours.get('reset', '')
 547
 548     colour_transitive = get_transitive_colourer(colours, len(vertices))
 549
 550     vspace = ' ' * vlen
 551     verticals = ''
 552     write("%*s %s  %sdestination%s" % (vlen, '',
 553                                        ' ' * len(vertices),
 554                                        c_header,
 555                                        c_reset))
 556     for i, v in enumerate(vertices):
 557         j = len(vertices) - i
 558         c = colour_cycle[i % len(colour_cycle)]
 559         if j == 1:
 560             start = '%s%ssource%s' % (vspace[:-6], c_header, c_reset)
 561         else:
 562             start = vspace
 563         write('%s %s%s%s%s%s %s%s' % (start,
 564                                       verticals,
 565                                       c_reset,
 566                                       c,
 567                                       corner,
 568                                       horizontal * j,
 569                                       v,
 570                                       c_reset
 571         ))
 572         verticals += c + vertical
 573
 574     connections = find_transitive_distance(vertices, edges)
 575
 576     for i, v in enumerate(vertices):
 577         c = colour_cycle[i % len(colour_cycle)]
 578         links = connections[v]
 579         row = []
 580         for v2 in vertices:
 581             link = links.get(v2)
 582             if link is None:
 583                 row.append('%s%s' % (c_disconn, missing))
 584                 continue
 585             if link == 0:
 586                 row.append('%s%s%s%s' % (c_reset, c, diagonal, c_reset))
 587             elif link == 1:
 588                 row.append('%s1%s' % (c_conn, c_reset))
 589             else:
 590                 ct = colour_transitive(link)
 591                 if link > 9:
 592                     link = '+'
 593                 row.append('%s%s%s' % (ct, link, c_reset))
 594
 595         write('%s%*s%s %s%s' % (c, vlen, v, c_reset,
 596                                 ''.join(row), c_reset))
 597
 598     if shorten_names:
 599         write('')
 600         for substitute, original in reversed(replacements):
 601             write("'%s%s%s' stands for '%s%s%s'" % (colour_cycle[0],
 602                                                     substitute,
 603                                                     c_reset,
 604                                                     colour_cycle[0],
 605                                                     original,
 606                                                     c_reset))
 607     if generate_key:
 608         write('')
 609         write("Data can get from %ssource%s to %sdestination%s in the "
 610               "indicated number of steps." % (c_header, c_reset,
 611                                               c_header, c_reset))
 612         write("%s%s%s means zero steps (it is the same DC)" %
 613               (colour_cycle[0], diagonal, c_reset))
 614         write("%s1%s means a direct link" % (c_conn, c_reset))
 615         write("%s2%s means a transitive link involving two steps "
 616               "(i.e. one intermediate DC)" %
 617               (colour_transitive(2), c_reset))
 618         write("%s%s%s means there is no connection, even through other DCs" %
 619               (c_disconn, missing, c_reset))
 620
 621     return '\n'.join(lines)