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
=[], 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.
156 for i
, item
in enumerate(key_items
):
157 is_vertex
, style
, label
= item
158 tag
= '%s%d_' % (prefix
, i
)
159 label
= quote_graph_label(label
)
160 name
= '%s_label' % tag
163 order_lines
.append(name
)
164 vertex_names
.append(name
)
165 vertex_lines
.append('%s[label="%s"; %s]' %
166 (name
, label
, style
))
168 edge_names
.append(name
)
171 order_lines
.append(name
)
172 edge_lines
.append('subgraph cluster_%s {' % tag
)
173 edge_lines
.append('%s[label=src; color="#000000"; group="%s_g"]' %
175 edge_lines
.append('%s[label=dest; color="#000000"; group="%s_g"]' %
177 edge_lines
.append('%s -> %s [constraint = false; %s]' % (e1
, e2
,
179 edge_lines
.append(('%s[shape=plaintext; style=solid; width=%f; '
181 (name
, width
, label
))
182 edge_lines
.append('}')
186 for i
, elision
in enumerate(reversed(elisions
)):
187 order_lines
.append('elision%d' % i
)
188 short
, long = elision
189 if short
[0] == ',' and long[0] == ',':
192 elision_str
+= ('\nelision%d[shape=plaintext; style=solid; '
193 'label="\“%s” means “%s”\\r"]\n'
194 % ((i
, short
, long)))
198 for n
in nodes_above
:
199 above_lines
.append('"%s" -> %s [style=invis]' %
202 s
= ('subgraph cluster_key {\n'
204 'subgraph cluster_key_nodes {\n'
209 'subgraph cluster_key_edges {\n'
218 '%s [style=invis; weight=9]'
220 % (';\n'.join(vertex_lines
),
221 '\n'.join(edge_lines
),
222 ' '.join(edge_names
),
224 ';\n'.join(above_lines
),
225 ' -> '.join(order_lines
),
231 def dot_graph(vertices
, edges
,
234 reformat_labels
=True,
243 vertex_clusters
=None):
244 """Generate a Graphviz representation of a list of vertices and edges.
246 :param vertices: list of vertex names (optional).
247 :param edges: list of (vertex, vertex) pairs
248 :param directed: bool: whether the graph is directed
249 :param title: optional title for the graph
250 :param reformat_labels: whether to wrap long vertex labels
251 :param vertex_colors: if not None, a sequence of colours for the vertices
252 :param edge_colors: if not None, colours for the edges
253 :param edge_labels: if not None, labels for the edges
254 :param vertex_styles: if not None, DOT style strings for vertices
255 :param edge_styles: if not None, DOT style strings for edges
256 :param graph_name: if not None, name of graph
257 :param shorten_names: if True, remove common DN suffixes
258 :param key: (is_vertex, style, description) tuples
259 :param vertex_clusters: list of subgraph cluster names
261 Colour, style, and label lists must be the same length as the
262 corresponding list of edges or vertices (or None).
264 Colours can be HTML RGB strings ("#FF0000") or common names
265 ("red"), or some other formats you don't want to think about.
267 If `vertices` is None, only the vertices mentioned in the edges
268 are shown, and their appearance can be modified using the
269 vertex_colors and vertex_styles arguments. Vertices appearing in
270 the edges but not in the `vertices` list will be shown but their
271 styles can not be modified.
277 vertices
= set(x
[0] for x
in edges
) |
set(x
[1] for x
in edges
)
280 vlist
= list(set(x
[0] for x
in edges
) |
281 set(x
[1] for x
in edges
) |
283 vmap
, elisions
= shorten_vertex_names(vlist
)
284 vertices
= [vmap
[x
] for x
in vertices
]
285 edges
= [(vmap
[a
], vmap
[b
]) for a
, b
in edges
]
290 if graph_name
is None:
291 graph_name
= 'A_samba_tool_production'
294 graph_type
= 'digraph'
300 write('/* generated by samba */')
301 write('%s %s {' % (graph_type
, graph_name
))
302 if title
is not None:
303 write('label="%s";' % (title
,))
304 write('fontsize=%s;\n' % (FONT_SIZE
))
305 write('node[fontname=Helvetica; fontsize=%s];\n' % (FONT_SIZE
))
310 for i
, v
in enumerate(vertices
):
311 v
= quote_graph_label(v
, reformat_labels
)
312 quoted_vertices
.append(v
)
314 if vertex_clusters
and vertex_clusters
[i
]:
315 cluster
= vertex_clusters
[i
]
316 if cluster
!= prev_cluster
:
317 if prev_cluster
is not None:
319 prev_cluster
= cluster
320 n
= quote_graph_label(cluster
)
322 write('subgraph cluster_%d {' % cluster_n
)
324 write('style = "rounded,dotted";')
325 write('node [style="filled"; fillcolor=white];')
326 write('label = "%s";' % n
)
328 if vertex_styles
and vertex_styles
[i
]:
329 attrs
.append(vertex_styles
[i
])
330 if vertex_colors
and vertex_colors
[i
]:
331 attrs
.append('color="%s"' % quote_graph_label(vertex_colors
[i
]))
333 write('"%s" [%s];' % (v
, ', '.join(attrs
)))
335 write('"%s";' % (v
,))
340 for i
, edge
in enumerate(edges
):
343 a
= "Missing source value"
345 b
= "Missing destination value"
347 a
= quote_graph_label(a
, reformat_labels
)
348 b
= quote_graph_label(b
, reformat_labels
)
352 label
= quote_graph_label(edge_labels
[i
])
353 attrs
.append('label="%s"' % label
)
355 attrs
.append('color="%s"' % quote_graph_label(edge_colors
[i
]))
357 attrs
.append(edge_styles
[i
]) # no quoting
359 write('"%s" %s "%s" [%s];' % (a
, connector
, b
, ', '.join(attrs
)))
361 write('"%s" %s "%s";' % (a
, connector
, b
))
364 key
= compile_graph_key(key_items
, nodes_above
=quoted_vertices
,
369 return '\n'.join(out
)
374 'alternate rows': (colour
.DARK_WHITE
, colour
.BLACK
),
375 'disconnected': colour
.RED
,
376 'connected': colour
.GREEN
,
377 'transitive': colour
.DARK_YELLOW
,
378 'header': colour
.UNDERLINE
,
379 'reset': colour
.C_NORMAL
,
382 'alternate rows': (colour
.DARK_WHITE
, colour
.BLACK
),
383 'disconnected': colour
.REV_RED
,
384 'connected': colour
.REV_GREEN
,
385 'transitive': colour
.REV_DARK_YELLOW
,
386 'header': colour
.UNDERLINE
,
387 'reset': colour
.C_NORMAL
,
390 'alternate rows': (colour
.xterm_256_colour(39),
391 colour
.xterm_256_colour(45)),
392 # 'alternate rows': (colour.xterm_256_colour(246),
393 # colour.xterm_256_colour(247)),
394 'disconnected': colour
.xterm_256_colour(124, bg
=True),
395 'connected': colour
.xterm_256_colour(112),
396 'transitive': colour
.xterm_256_colour(214),
397 'transitive scale': (colour
.xterm_256_colour(190),
398 colour
.xterm_256_colour(184),
399 colour
.xterm_256_colour(220),
400 colour
.xterm_256_colour(214),
401 colour
.xterm_256_colour(208),
403 'header': colour
.UNDERLINE
,
404 'reset': colour
.C_NORMAL
,
406 'xterm-256color-heatmap': {
407 'alternate rows': (colour
.xterm_256_colour(171),
408 colour
.xterm_256_colour(207)),
409 # 'alternate rows': (colour.xterm_256_colour(246),
410 # colour.xterm_256_colour(247)),
411 'disconnected': colour
.xterm_256_colour(124, bg
=True),
412 'connected': colour
.xterm_256_colour(112, bg
=True),
413 'transitive': colour
.xterm_256_colour(214, bg
=True),
414 'transitive scale': (colour
.xterm_256_colour(190, bg
=True),
415 colour
.xterm_256_colour(184, bg
=True),
416 colour
.xterm_256_colour(220, bg
=True),
417 colour
.xterm_256_colour(214, bg
=True),
418 colour
.xterm_256_colour(208, bg
=True),
420 'header': colour
.UNDERLINE
,
421 'reset': colour
.C_NORMAL
,
424 'alternate rows': ('',),
455 def find_transitive_distance(vertices
, edges
):
456 all_vertices
= (set(vertices
) |
457 set(e
[0] for e
in edges
) |
458 set(e
[1] for e
in edges
))
460 if all_vertices
!= set(vertices
):
461 print("there are unknown vertices: %s" %
462 (all_vertices
- set(vertices
)),
465 # with n vertices, we are always less than n hops away from
467 inf
= len(all_vertices
)
469 for v
in all_vertices
:
470 distances
[v
] = {v
: 0}
472 for src
, dest
in edges
:
473 distances
[src
][dest
] = distances
[src
].get(dest
, 1)
475 # This algorithm (and implementation) seems very suboptimal.
476 # potentially O(n^4), though n is smallish.
480 for v
, d
in distances
.items():
482 new_distances
[v
] = new_d
483 for dest
, cost
in d
.items():
484 for leaf
, cost2
in distances
[dest
].items():
485 new_cost
= cost
+ cost2
486 old_cost
= d
.get(leaf
, inf
)
487 if new_cost
< old_cost
:
488 new_d
[leaf
] = new_cost
491 distances
= new_distances
495 # filter out unwanted vertices and infinite links
500 a
= distances
[v
].get(v2
, inf
)
507 def get_transitive_colourer(colours
, n_vertices
):
508 if 'transitive scale' in colours
:
509 scale
= colours
['transitive scale']
511 n
= 1 + int(n_vertices
** 0.5)
514 if not isinstance(link
, int):
516 return scale
[min(link
* m
// n
, m
- 1)]
520 return colours
['transitive']
525 def distance_matrix(vertices
, edges
,
530 grouping_function
=None,
535 charset
= CHARSETS
['utf8' if utf8
else 'ascii']
536 vertical
= charset
['vertical']
537 horizontal
= charset
['horizontal']
538 corner
= charset
['corner']
539 diagonal
= charset
['diagonal']
540 missing
= charset
['missing']
541 right_arrow
= charset
['right_arrow']
543 colours
= COLOUR_SETS
[colour
]
545 colour_cycle
= cycle(colours
.get('alternate rows', ('',)))
548 vertices
= sorted(set(x
[0] for x
in edges
) |
set(x
[1] for x
in edges
))
550 if grouping_function
is not None:
551 # we sort and colour according to the grouping function
552 # which can be used to e.g. alternate colours by site.
553 vertices
= sorted(vertices
, key
=grouping_function
)
555 for k
, v
in groupby(vertices
, key
=grouping_function
):
556 c
= next(colour_cycle
)
557 colour_list
.extend(c
for x
in v
)
559 colour_list
= [next(colour_cycle
) for v
in vertices
]
562 vlist
= list(set(x
[0] for x
in edges
) |
563 set(x
[1] for x
in edges
) |
565 vmap
, replacements
= shorten_vertex_names(vlist
, '+',
567 vertices
= [vmap
[x
] for x
in vertices
]
568 edges
= [(vmap
[a
], vmap
[b
]) for a
, b
in edges
]
570 vlen
= max(6, max(len(v
) for v
in vertices
))
572 # first, the key for the columns
573 c_header
= colours
.get('header', '')
574 c_disconn
= colours
.get('disconnected', '')
575 c_conn
= colours
.get('connected', '')
576 c_reset
= colours
.get('reset', '')
578 colour_transitive
= get_transitive_colourer(colours
, len(vertices
))
582 write("%*s %s %sdestination%s" % (vlen
, '',
586 for i
, v
in enumerate(vertices
):
587 j
= len(vertices
) - i
590 start
= '%s%ssource%s' % (vspace
[:-6], c_header
, c_reset
)
593 write('%s %s%s%s%s%s %s%s' % (start
,
602 verticals
+= c
+ vertical
604 connections
= find_transitive_distance(vertices
, edges
)
606 for i
, v
in enumerate(vertices
):
608 links
= connections
[v
]
613 row
.append('%s%s' % (c_disconn
, missing
))
616 row
.append('%s%s%s%s' % (c_reset
, c
, diagonal
, c_reset
))
618 row
.append('%s1%s' % (c_conn
, c_reset
))
620 ct
= colour_transitive(link
)
623 row
.append('%s%s%s' % (ct
, link
, c_reset
))
625 if row_comments
is not None and row_comments
[i
]:
626 row
.append('%s %s %s' % (c_reset
, right_arrow
, row_comments
[i
]))
628 write('%s%*s%s %s%s' % (c
, vlen
, v
, c_reset
,
629 ''.join(row
), c_reset
))
631 example_c
= next(colour_cycle
)
634 for substitute
, original
in reversed(replacements
):
635 write("'%s%s%s' stands for '%s%s%s'" % (example_c
,
643 write("Data can get from %ssource%s to %sdestination%s in the "
644 "indicated number of steps." % (c_header
, c_reset
,
646 write("%s%s%s means zero steps (it is the same DC)" %
647 (example_c
, diagonal
, c_reset
))
648 write("%s1%s means a direct link" % (c_conn
, c_reset
))
649 write("%s2%s means a transitive link involving two steps "
650 "(i.e. one intermediate DC)" %
651 (colour_transitive(2), c_reset
))
652 write("%s%s%s means there is no connection, even through other DCs" %
653 (c_disconn
, missing
, c_reset
))
655 return '\n'.join(lines
)
658 def pad_char(char
, digits
, padding
=' '):
661 return ' ' * (digits
- 1) + char
+ padding
664 def transpose_dict_matrix(m
):
666 for k1
, row
in m
.items():
667 for k2
, dist
in row
.items():
668 m2
.setdefault(k2
, {})[k1
] = dist
672 def full_matrix(rows
,
677 grouping_function
=None,
682 xlabel
='destination',
688 rows
= transpose_dict_matrix(rows
)
690 use_padding
= digits
> 1
692 charset
= CHARSETS
['utf8' if utf8
else 'ascii']
693 vertical
= pad_char(charset
['vertical'], digits
)
694 horizontal
= charset
['horizontal'] * (digits
+ use_padding
)
695 corner
= pad_char(charset
['corner'], digits
,
696 charset
['horizontal'])
697 diagonal
= pad_char(charset
['diagonal'], digits
)
698 missing
= pad_char(charset
['missing'], digits
)
699 toobig
= pad_char('>', digits
)
700 right_arrow
= charset
['right_arrow']
701 empty
= pad_char(' ', digits
)
703 colours
= COLOUR_SETS
[colour
]
705 colour_cycle
= cycle(colours
.get('alternate rows', ('',)))
706 vertices
= list(rows
.keys())
707 if grouping_function
is not None:
708 # we sort and colour according to the grouping function
709 # which can be used to e.g. alternate colours by site.
710 vertices
.sort(key
=grouping_function
)
712 for k
, v
in groupby(vertices
, key
=grouping_function
):
713 c
= next(colour_cycle
)
714 colour_list
.extend(c
for x
in v
)
716 colour_list
= [next(colour_cycle
) for v
in vertices
]
719 vmap
, replacements
= shorten_vertex_names(vertices
, '+',
722 for vert
, r
in rows
.items():
723 rows2
[vmap
[vert
]] = dict((vmap
[k
], v
) for k
, v
in r
.items())
726 vertices
= list(rows
.keys())
728 vlen
= max(6, len(xlabel
), max(len(v
) for v
in vertices
))
730 # first, the key for the columns
731 c_header
= colours
.get('header', '')
732 c_disconn
= colours
.get('disconnected', '')
733 c_conn
= colours
.get('connected', '')
734 c_reset
= colours
.get('reset', '')
736 if colour_scale
is None:
737 colour_scale
= len(rows
)
738 colour_transitive
= get_transitive_colourer(colours
, colour_scale
)
742 write("%s %s %s%s%s" % (vspace
,
743 empty
* (len(rows
) + 1),
747 for i
, v
in enumerate(vertices
):
751 start
= '%s%s%s%s' % (vspace
[:-len(ylabel
)],
757 write('%s %s%s%s%s%s %s%s' % (start
,
766 verticals
+= '%s%s' % (c
, vertical
)
768 end_cell
= '%s%s' % (' ' * use_padding
, c_reset
)
770 for i
, v
in enumerate(vertices
):
776 row
.append('%s%s%s' % (c_disconn
, missing
, c_reset
))
778 row
.append('%s%s%s%s' % (c_reset
, c
, diagonal
, c_reset
))
781 if link
>= 10 ** digits
:
782 ct
= colour_transitive(link
)
783 row
.append('%s%s%s' % (ct
, toobig
, c_reset
))
789 ct
= colour_transitive(link
)
790 row
.append('%s%*s%s' % (ct
, digits
, link
, end_cell
))
792 if row_comments
is not None and row_comments
[i
]:
793 row
.append('%s %s %s' % (c_reset
, right_arrow
, row_comments
[i
]))
795 write('%s%*s%s %s%s' % (c
, vlen
, v
, c_reset
,
796 ''.join(row
), c_reset
))
798 if overflow
or shorten_names
:
802 write("'%s%s%s' means greater than %d " %
803 (colour_transitive(10 ** digits
),
809 example_c
= next(colour_cycle
)
810 for substitute
, original
in reversed(replacements
):
811 write("'%s%s%s' stands for '%s%s%s'" % (example_c
,
818 return '\n'.join(lines
)