1 # -*- coding: utf-8 -*-
2 # Graph topology utilities and dot file generation
4 # Copyright (C) Andrew Bartlett 2018.
6 # Written by Douglas Bagnall <douglas.bagnall@catalyst.net.nz>
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.
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.
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/>.
21 from samba
import colour
23 from itertools
import cycle
, groupby
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."""
35 s
= s
.replace(',', ',\n')
37 for p
in s
.split('\n'):
40 q
, p
= p
.split('-', 1)
45 pieces
.append(q
+ '-')
49 return '\\n'.join(pieces
)
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('"', '\"')
58 s
= reformat_graph_label(s
)
62 def shorten_vertex_names(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.
68 :param vertices: a sequence of vertices to shorten
69 :param suffix: the replacement string [",..."]
70 :param aggressive: replace certain common non-suffix strings
72 :return: tuple of (rename map, replacements)
74 The rename map is a dictionary mapping the old vertex names to
75 their shortened versions. If no changes are made, replacements
78 vmap
= dict((v
, v
) for v
in vertices
)
82 # walk backwards along all the strings until we meet a character
83 # that is not shared by all.
85 vlist
= list(vmap
.values())
88 c
= set(x
[i
] for x
in vlist
)
89 if len(c
) > 1 or '*' in c
:
93 # We have indexed beyond the start of a string, which should
94 # only happen if one node is a strict suffix of all others.
95 return vmap
, replacements
97 # add one to get to the last unanimous character.
100 # now, we actually really want to split on a comma. So we walk
103 while i
< len(x
) and x
[i
] != ',':
106 if i
>= -len(suffix
):
107 # there is nothing to gain here
108 return vmap
, replacements
110 replacements
.append((suffix
, x
[i
:]))
112 for k
, v
in vmap
.items():
113 vmap
[k
] = v
[:i
] + suffix
116 # Remove known common annoying strings
117 for v
in vmap
.values():
118 if ',CN=Servers,' not in v
:
121 vmap
= dict((k
, v
.replace(',CN=Servers,', ',**,', 1))
122 for k
, v
in vmap
.items())
123 replacements
.append(('**', 'CN=Servers'))
125 for v
in vmap
.values():
126 if not v
.startswith('CN=NTDS Settings,'):
129 vmap
= dict((k
, v
.replace('CN=NTDS Settings,', '*,', 1))
130 for k
, v
in vmap
.items())
131 replacements
.append(('*', 'CN=NTDS Settings'))
133 return vmap
, replacements
136 def compile_graph_key(key_items
, nodes_above
=None, elisions
=None,
137 prefix
='key_', width
=2):
138 """Generate a dot file snippet that acts as a legend for a graph.
140 :param key_items: sequence of items (is_vertex, style, label)
141 :param nodes_above: list of vertices (pushes key into right position)
142 :param elision: tuple (short, full) indicating suffix replacement
143 :param prefix: string used to generate key node names ["key_"]
144 :param width: default width of node lines
146 Each item in key_items is a tuple of (is_vertex, style, label).
147 is_vertex is a boolean indicating whether the item is a vertex
148 (True) or edge (False). Style is a dot style string for the edge
149 or vertex. label is the text associated with the key item.
151 if nodes_above
is None:
158 for i
, item
in enumerate(key_items
):
159 is_vertex
, style
, label
= item
160 tag
= '%s%d_' % (prefix
, i
)
161 label
= quote_graph_label(label
)
162 name
= '%s_label' % tag
165 order_lines
.append(name
)
166 vertex_names
.append(name
)
167 vertex_lines
.append('%s[label="%s"; %s]' %
168 (name
, label
, style
))
170 edge_names
.append(name
)
173 order_lines
.append(name
)
174 edge_lines
.append('subgraph cluster_%s {' % tag
)
175 edge_lines
.append('%s[label=src; color="#000000"; group="%s_g"]' %
177 edge_lines
.append('%s[label=dest; color="#000000"; group="%s_g"]' %
179 edge_lines
.append('%s -> %s [constraint = false; %s]' % (e1
, e2
,
181 edge_lines
.append(('%s[shape=plaintext; style=solid; width=%f; '
183 (name
, width
, label
))
184 edge_lines
.append('}')
188 for i
, elision
in enumerate(reversed(elisions
)):
189 order_lines
.append('elision%d' % i
)
190 short
, long = elision
191 if short
[0] == ',' and long[0] == ',':
194 elision_str
+= ('\nelision%d[shape=plaintext; style=solid; '
195 'label="\\“%s” means “%s”\\r"]\n'
196 % ((i
, short
, long)))
200 for n
in nodes_above
:
201 above_lines
.append('"%s" -> %s [style=invis]' %
204 s
= ('subgraph cluster_key {\n'
206 'subgraph cluster_key_nodes {\n'
211 'subgraph cluster_key_edges {\n'
220 '%s [style=invis; weight=9]'
222 % (';\n'.join(vertex_lines
),
223 '\n'.join(edge_lines
),
224 ' '.join(edge_names
),
226 ';\n'.join(above_lines
),
227 ' -> '.join(order_lines
),
233 def dot_graph(vertices
, edges
,
236 reformat_labels
=True,
245 vertex_clusters
=None):
246 """Generate a Graphviz representation of a list of vertices and edges.
248 :param vertices: list of vertex names (optional).
249 :param edges: list of (vertex, vertex) pairs
250 :param directed: bool: whether the graph is directed
251 :param title: optional title for the graph
252 :param reformat_labels: whether to wrap long vertex labels
253 :param vertex_colors: if not None, a sequence of colours for the vertices
254 :param edge_colors: if not None, colours for the edges
255 :param edge_labels: if not None, labels for the edges
256 :param vertex_styles: if not None, DOT style strings for vertices
257 :param edge_styles: if not None, DOT style strings for edges
258 :param graph_name: if not None, name of graph
259 :param shorten_names: if True, remove common DN suffixes
260 :param key: (is_vertex, style, description) tuples
261 :param vertex_clusters: list of subgraph cluster names
263 Colour, style, and label lists must be the same length as the
264 corresponding list of edges or vertices (or None).
266 Colours can be HTML RGB strings ("#FF0000") or common names
267 ("red"), or some other formats you don't want to think about.
269 If `vertices` is None, only the vertices mentioned in the edges
270 are shown, and their appearance can be modified using the
271 vertex_colors and vertex_styles arguments. Vertices appearing in
272 the edges but not in the `vertices` list will be shown but their
273 styles can not be modified.
279 vertices
= set(x
[0] for x
in edges
) |
set(x
[1] for x
in edges
)
282 vlist
= list(set(x
[0] for x
in edges
) |
283 set(x
[1] for x
in edges
) |
285 vmap
, elisions
= shorten_vertex_names(vlist
)
286 vertices
= [vmap
[x
] for x
in vertices
]
287 edges
= [(vmap
[a
], vmap
[b
]) for a
, b
in edges
]
292 if graph_name
is None:
293 graph_name
= 'A_samba_tool_production'
296 graph_type
= 'digraph'
302 write('/* generated by samba */')
303 write('%s %s {' % (graph_type
, graph_name
))
304 if title
is not None:
305 write('label="%s";' % (title
,))
306 write('fontsize=%s;\n' % (FONT_SIZE
))
307 write('node[fontname=Helvetica; fontsize=%s];\n' % (FONT_SIZE
))
312 for i
, v
in enumerate(vertices
):
313 v
= quote_graph_label(v
, reformat_labels
)
314 quoted_vertices
.append(v
)
316 if vertex_clusters
and vertex_clusters
[i
]:
317 cluster
= vertex_clusters
[i
]
318 if cluster
!= prev_cluster
:
319 if prev_cluster
is not None:
321 prev_cluster
= cluster
322 n
= quote_graph_label(cluster
)
324 write('subgraph cluster_%d {' % cluster_n
)
326 write('style = "rounded,dotted";')
327 write('node [style="filled"; fillcolor=white];')
328 write('label = "%s";' % n
)
330 if vertex_styles
and vertex_styles
[i
]:
331 attrs
.append(vertex_styles
[i
])
332 if vertex_colors
and vertex_colors
[i
]:
333 attrs
.append('color="%s"' % quote_graph_label(vertex_colors
[i
]))
335 write('"%s" [%s];' % (v
, ', '.join(attrs
)))
337 write('"%s";' % (v
,))
342 for i
, edge
in enumerate(edges
):
345 a
= "Missing source value"
347 b
= "Missing destination value"
349 a
= quote_graph_label(a
, reformat_labels
)
350 b
= quote_graph_label(b
, reformat_labels
)
354 label
= quote_graph_label(edge_labels
[i
])
355 attrs
.append('label="%s"' % label
)
357 attrs
.append('color="%s"' % quote_graph_label(edge_colors
[i
]))
359 attrs
.append(edge_styles
[i
]) # no quoting
361 write('"%s" %s "%s" [%s];' % (a
, connector
, b
, ', '.join(attrs
)))
363 write('"%s" %s "%s";' % (a
, connector
, b
))
366 key
= compile_graph_key(key_items
, nodes_above
=quoted_vertices
,
371 return '\n'.join(out
)
376 'alternate rows': (colour
.DARK_WHITE
, colour
.BLACK
),
377 'disconnected': colour
.RED
,
378 'connected': colour
.GREEN
,
379 'transitive': colour
.DARK_YELLOW
,
380 'header': colour
.UNDERLINE
,
381 'reset': colour
.C_NORMAL
,
384 'alternate rows': (colour
.DARK_WHITE
, colour
.BLACK
),
385 'disconnected': colour
.REV_RED
,
386 'connected': colour
.REV_GREEN
,
387 'transitive': colour
.REV_DARK_YELLOW
,
388 'header': colour
.UNDERLINE
,
389 'reset': colour
.C_NORMAL
,
392 'alternate rows': (colour
.xterm_256_colour(39),
393 colour
.xterm_256_colour(45)),
394 # 'alternate rows': (colour.xterm_256_colour(246),
395 # colour.xterm_256_colour(247)),
396 'disconnected': colour
.xterm_256_colour(124, bg
=True),
397 'connected': colour
.xterm_256_colour(112),
398 'transitive': colour
.xterm_256_colour(214),
399 'transitive scale': (colour
.xterm_256_colour(190),
400 colour
.xterm_256_colour(184),
401 colour
.xterm_256_colour(220),
402 colour
.xterm_256_colour(214),
403 colour
.xterm_256_colour(208),
405 'header': colour
.UNDERLINE
,
406 'reset': colour
.C_NORMAL
,
408 'xterm-256color-heatmap': {
409 'alternate rows': (colour
.xterm_256_colour(171),
410 colour
.xterm_256_colour(207)),
411 # 'alternate rows': (colour.xterm_256_colour(246),
412 # colour.xterm_256_colour(247)),
413 'disconnected': colour
.xterm_256_colour(124, bg
=True),
414 'connected': colour
.xterm_256_colour(112, bg
=True),
415 'transitive': colour
.xterm_256_colour(214, bg
=True),
416 'transitive scale': (colour
.xterm_256_colour(190, bg
=True),
417 colour
.xterm_256_colour(184, bg
=True),
418 colour
.xterm_256_colour(220, bg
=True),
419 colour
.xterm_256_colour(214, bg
=True),
420 colour
.xterm_256_colour(208, bg
=True),
422 'header': colour
.UNDERLINE
,
423 'reset': colour
.C_NORMAL
,
426 'alternate rows': ('',),
457 def find_transitive_distance(vertices
, edges
):
458 all_vertices
= (set(vertices
) |
459 set(e
[0] for e
in edges
) |
460 set(e
[1] for e
in edges
))
462 if all_vertices
!= set(vertices
):
463 print("there are unknown vertices: %s" %
464 (all_vertices
- set(vertices
)),
467 # with n vertices, we are always less than n hops away from
469 inf
= len(all_vertices
)
471 for v
in all_vertices
:
472 distances
[v
] = {v
: 0}
474 for src
, dest
in edges
:
475 distances
[src
][dest
] = distances
[src
].get(dest
, 1)
477 # This algorithm (and implementation) seems very suboptimal.
478 # potentially O(n^4), though n is smallish.
482 for v
, d
in distances
.items():
484 new_distances
[v
] = new_d
485 for dest
, cost
in d
.items():
486 for leaf
, cost2
in distances
[dest
].items():
487 new_cost
= cost
+ cost2
488 old_cost
= d
.get(leaf
, inf
)
489 if new_cost
< old_cost
:
490 new_d
[leaf
] = new_cost
493 distances
= new_distances
497 # filter out unwanted vertices and infinite links
502 a
= distances
[v
].get(v2
, inf
)
509 def get_transitive_colourer(colours
, n_vertices
):
510 if 'transitive scale' in colours
:
511 scale
= colours
['transitive scale']
513 n
= 1 + int(n_vertices
** 0.5)
516 if not isinstance(link
, int):
518 return scale
[min(link
* m
// n
, m
- 1)]
522 return colours
['transitive']
527 def distance_matrix(vertices
, edges
,
532 grouping_function
=None,
537 charset
= CHARSETS
['utf8' if utf8
else 'ascii']
538 vertical
= charset
['vertical']
539 horizontal
= charset
['horizontal']
540 corner
= charset
['corner']
541 diagonal
= charset
['diagonal']
542 missing
= charset
['missing']
543 right_arrow
= charset
['right_arrow']
545 colours
= COLOUR_SETS
[colour
]
547 colour_cycle
= cycle(colours
.get('alternate rows', ('',)))
550 vertices
= sorted(set(x
[0] for x
in edges
) |
set(x
[1] for x
in edges
))
552 if grouping_function
is not None:
553 # we sort and colour according to the grouping function
554 # which can be used to e.g. alternate colours by site.
555 vertices
= sorted(vertices
, key
=grouping_function
)
557 for k
, v
in groupby(vertices
, key
=grouping_function
):
558 c
= next(colour_cycle
)
559 colour_list
.extend(c
for x
in v
)
561 colour_list
= [next(colour_cycle
) for v
in vertices
]
564 vlist
= list(set(x
[0] for x
in edges
) |
565 set(x
[1] for x
in edges
) |
567 vmap
, replacements
= shorten_vertex_names(vlist
, '+',
569 vertices
= [vmap
[x
] for x
in vertices
]
570 edges
= [(vmap
[a
], vmap
[b
]) for a
, b
in edges
]
572 vlen
= max(6, max(len(v
) for v
in vertices
))
574 # first, the key for the columns
575 c_header
= colours
.get('header', '')
576 c_disconn
= colours
.get('disconnected', '')
577 c_conn
= colours
.get('connected', '')
578 c_reset
= colours
.get('reset', '')
580 colour_transitive
= get_transitive_colourer(colours
, len(vertices
))
584 write("%*s %s %sdestination%s" % (vlen
, '',
588 for i
, v
in enumerate(vertices
):
589 j
= len(vertices
) - i
592 start
= '%s%ssource%s' % (vspace
[:-6], c_header
, c_reset
)
595 write('%s %s%s%s%s%s %s%s' % (start
,
604 verticals
+= c
+ vertical
606 connections
= find_transitive_distance(vertices
, edges
)
608 for i
, v
in enumerate(vertices
):
610 links
= connections
[v
]
615 row
.append('%s%s' % (c_disconn
, missing
))
618 row
.append('%s%s%s%s' % (c_reset
, c
, diagonal
, c_reset
))
620 row
.append('%s1%s' % (c_conn
, c_reset
))
622 ct
= colour_transitive(link
)
625 row
.append('%s%s%s' % (ct
, link
, c_reset
))
627 if row_comments
is not None and row_comments
[i
]:
628 row
.append('%s %s %s' % (c_reset
, right_arrow
, row_comments
[i
]))
630 write('%s%*s%s %s%s' % (c
, vlen
, v
, c_reset
,
631 ''.join(row
), c_reset
))
633 example_c
= next(colour_cycle
)
636 for substitute
, original
in reversed(replacements
):
637 write("'%s%s%s' stands for '%s%s%s'" % (example_c
,
645 write("Data can get from %ssource%s to %sdestination%s in the "
646 "indicated number of steps." % (c_header
, c_reset
,
648 write("%s%s%s means zero steps (it is the same DC)" %
649 (example_c
, diagonal
, c_reset
))
650 write("%s1%s means a direct link" % (c_conn
, c_reset
))
651 write("%s2%s means a transitive link involving two steps "
652 "(i.e. one intermediate DC)" %
653 (colour_transitive(2), c_reset
))
654 write("%s%s%s means there is no connection, even through other DCs" %
655 (c_disconn
, missing
, c_reset
))
657 return '\n'.join(lines
)
660 def pad_char(char
, digits
, padding
=' '):
663 return ' ' * (digits
- 1) + char
+ padding
666 def transpose_dict_matrix(m
):
668 for k1
, row
in m
.items():
669 for k2
, dist
in row
.items():
670 m2
.setdefault(k2
, {})[k1
] = dist
674 def full_matrix(rows
,
679 grouping_function
=None,
684 xlabel
='destination',
690 rows
= transpose_dict_matrix(rows
)
692 use_padding
= digits
> 1
694 charset
= CHARSETS
['utf8' if utf8
else 'ascii']
695 vertical
= pad_char(charset
['vertical'], digits
)
696 horizontal
= charset
['horizontal'] * (digits
+ use_padding
)
697 corner
= pad_char(charset
['corner'], digits
,
698 charset
['horizontal'])
699 diagonal
= pad_char(charset
['diagonal'], digits
)
700 missing
= pad_char(charset
['missing'], digits
)
701 toobig
= pad_char('>', digits
)
702 right_arrow
= charset
['right_arrow']
703 empty
= pad_char(' ', digits
)
705 colours
= COLOUR_SETS
[colour
]
707 colour_cycle
= cycle(colours
.get('alternate rows', ('',)))
708 vertices
= list(rows
.keys())
709 if grouping_function
is not None:
710 # we sort and colour according to the grouping function
711 # which can be used to e.g. alternate colours by site.
712 vertices
.sort(key
=grouping_function
)
714 for k
, v
in groupby(vertices
, key
=grouping_function
):
715 c
= next(colour_cycle
)
716 colour_list
.extend(c
for x
in v
)
718 colour_list
= [next(colour_cycle
) for v
in vertices
]
721 vmap
, replacements
= shorten_vertex_names(vertices
, '+',
724 for vert
, r
in rows
.items():
725 rows2
[vmap
[vert
]] = dict((vmap
[k
], v
) for k
, v
in r
.items())
728 vertices
= list(rows
.keys())
730 vlen
= max(6, len(xlabel
), max(len(v
) for v
in vertices
))
732 # first, the key for the columns
733 c_header
= colours
.get('header', '')
734 c_disconn
= colours
.get('disconnected', '')
735 c_conn
= colours
.get('connected', '')
736 c_reset
= colours
.get('reset', '')
738 if colour_scale
is None:
739 colour_scale
= len(rows
)
740 colour_transitive
= get_transitive_colourer(colours
, colour_scale
)
744 write("%s %s %s%s%s" % (vspace
,
745 empty
* (len(rows
) + 1),
749 for i
, v
in enumerate(vertices
):
753 start
= '%s%s%s%s' % (vspace
[:-len(ylabel
)],
759 write('%s %s%s%s%s%s %s%s' % (start
,
768 verticals
+= '%s%s' % (c
, vertical
)
770 end_cell
= '%s%s' % (' ' * use_padding
, c_reset
)
772 for i
, v
in enumerate(vertices
):
778 row
.append('%s%s%s' % (c_disconn
, missing
, c_reset
))
780 row
.append('%s%s%s%s' % (c_reset
, c
, diagonal
, c_reset
))
783 if link
>= 10 ** digits
:
784 ct
= colour_transitive(link
)
785 row
.append('%s%s%s' % (ct
, toobig
, c_reset
))
791 ct
= colour_transitive(link
)
792 row
.append('%s%*s%s' % (ct
, digits
, link
, end_cell
))
794 if row_comments
is not None and row_comments
[i
]:
795 row
.append('%s %s %s' % (c_reset
, right_arrow
, row_comments
[i
]))
797 write('%s%*s%s %s%s' % (c
, vlen
, v
, c_reset
,
798 ''.join(row
), c_reset
))
800 if overflow
or shorten_names
:
804 write("'%s%s%s' means greater than %d " %
805 (colour_transitive(10 ** digits
),
811 example_c
= next(colour_cycle
)
812 for substitute
, original
in reversed(replacements
):
813 write("'%s%s%s' stands for '%s%s%s'" % (example_c
,
820 return '\n'.join(lines
)