6cdd2ef77ed67c020234345b116f5ac6f0fc8cd8
[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 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.items())
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.items())
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.items():
467             new_d = d.copy()
468             new_distances[v] = new_d
469             for dest, cost in d.items():
470                 for leaf, cost2 in distances[dest].items():
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)