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