samba-mirror/python/samba/graph.py

# -*- coding: utf-8 -*-
# Graph topology utilities and dot file generation
#
# Copyright (C) Andrew Bartlett 2018.
#
# Written by Douglas Bagnall <douglas.bagnall@catalyst.net.nz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import print_function
from __future__ import division
from samba import colour
import sys
from itertools import cycle, groupby

FONT_SIZE = 10


def reformat_graph_label(s):
    """Break DNs over multiple lines, for better shaped and arguably more
    readable nodes. We try to split after commas, and if necessary
    after hyphens or failing that in arbitrary places."""
    if len(s) < 12:
        return s

    s = s.replace(',', ',\n')
    pieces = []
    for p in s.split('\n'):
        while len(p) > 20:
            if '-' in p[2:20]:
                q, p = p.split('-', 1)
            else:
                n = len(p) // 12
                b = len(p) // n
                q, p = p[:b], p[b:]
            pieces.append(q + '-')
        if p:
            pieces.append(p)

    return '\\n'.join(pieces)


def quote_graph_label(s, reformat=False):
    """Escape a string as graphvis requires."""
    # escaping inside quotes is simple in dot, because only " is escaped.
    # there is no need to count backslashes in sequences like \\\\"
    s = s.replace('"', '\"')
    if reformat:
        s = reformat_graph_label(s)
    return "%s" % s


def shorten_vertex_names(vertices, suffix=',...', aggressive=False):
    """Replace the common suffix (in practice, the base DN) of a number of
    vertices with a short string (default ",..."). If this seems
    pointless because the replaced string is very short or the results
    seem strange, the original vertices are retained.

    :param vertices: a sequence of vertices to shorten
    :param suffix: the replacement string [",..."]
    :param aggressive: replace certain common non-suffix strings

    :return: tuple of (rename map, replacements)

    The rename map is a dictionary mapping the old vertex names to
    their shortened versions. If no changes are made, replacements
    will be empty.
    """
    vmap = dict((v, v) for v in vertices)
    replacements = []

    if len(vmap) > 1:
        # walk backwards along all the strings until we meet a character
        # that is not shared by all.
        i = -1
        vlist = list(vmap.values())
        try:
            while True:
                c = set(x[i] for x in vlist)
                if len(c) > 1 or c == {'*'}:
                    break
                i -= 1
        except IndexError:
            # We have indexed beyond the start of a string, which should
            # only happen if one node is a strict suffix of all others.
            return vmap, replacements

        # add one to get to the last unanimous character.
        i += 1

        # now, we actually really want to split on a comma. So we walk
        # back to a comma.
        x = vlist[0]
        while i < len(x) and x[i] != ',':
            i += 1

        if i >= -len(suffix):
            # there is nothing to gain here
            return vmap, replacements

        replacements.append((suffix, x[i:]))

        for k, v in vmap.items():
            vmap[k] = v[:i] + suffix

    if aggressive:
        # Remove known common annoying strings
        for v in vmap.values():
            if ',CN=Servers,' not in v:
                break
        else:
            vmap = dict((k, v.replace(',CN=Servers,', ',**,', 1))
                       for k, v in vmap.items())
            replacements.append(('**', 'CN=Servers'))

        for v in vmap.values():
            if not v.startswith('CN=NTDS Settings,'):
                break
        else:
            vmap = dict((k, v.replace('CN=NTDS Settings,', '*,', 1))
                       for k, v in vmap.items())
            replacements.append(('*', 'CN=NTDS Settings'))

    return vmap, replacements






def compile_graph_key(key_items, nodes_above=[], elisions=None,
                      prefix='key_', width=2):
    """Generate a dot file snippet that acts as a legend for a graph.

    :param key_items: sequence of items (is_vertex, style, label)
    :param nodes_above: list of vertices (pushes key into right position)
    :param elision: tuple (short, full) indicating suffix replacement
    :param prefix: string used to generate key node names ["key_"]
    :param width: default width of node lines

    Each item in key_items is a tuple of (is_vertex, style, label).
    is_vertex is a boolean indicating whether the item is a vertex
    (True) or edge (False). Style is a dot style string for the edge
    or vertex. label is the text associated with the key item.
    """
    edge_lines = []
    edge_names = []
    vertex_lines = []
    vertex_names = []
    order_lines = []
    for i, item in enumerate(key_items):
        is_vertex, style, label = item
        tag = '%s%d_' % (prefix, i)
        label = quote_graph_label(label)
        name = '%s_label' % tag

        if is_vertex:
            order_lines.append(name)
            vertex_names.append(name)
            vertex_lines.append('%s[label="%s"; %s]' %
                                (name, label, style))
        else:
            edge_names.append(name)
            e1 = '%se1' % tag
            e2 = '%se2' % tag
            order_lines.append(name)
            edge_lines.append('subgraph cluster_%s {' % tag)
            edge_lines.append('%s[label=src; color="#000000"; group="%s_g"]' %
                              (e1, tag))
            edge_lines.append('%s[label=dest; color="#000000"; group="%s_g"]' %
                              (e2, tag))
            edge_lines.append('%s -> %s [constraint = false; %s]' % (e1, e2,
                                                                     style))
            edge_lines.append(('%s[shape=plaintext; style=solid; width=%f; '
                               'label="%s\\r"]') %
                              (name, width, label))
            edge_lines.append('}')

    elision_str = ''
    if elisions:
        for i, elision in enumerate(reversed(elisions)):
            order_lines.append('elision%d' % i)
            short, long = elision
            if short[0] == ',' and long[0] == ',':
                short = short[1:]
                long = long[1:]
            elision_str += ('\nelision%d[shape=plaintext; style=solid; '
                            'label="\“%s”  means  “%s”\\r"]\n'
                            % ((i, short, long)))

    above_lines = []
    if order_lines:
        for n in nodes_above:
            above_lines.append('"%s" -> %s [style=invis]' %
                               (n, order_lines[0]))

    s = ('subgraph cluster_key {\n'
         'label="Key";\n'
         'subgraph cluster_key_nodes {\n'
         'label="";\n'
         'color = "invis";\n'
         '%s\n'
         '}\n'
         'subgraph cluster_key_edges {\n'
         'label="";\n'
         'color = "invis";\n'
         '%s\n'
         '{%s}\n'
         '}\n'
         '%s\n'
         '}\n'
         '%s\n'
         '%s [style=invis; weight=9]'
         '\n'
         % (';\n'.join(vertex_lines),
            '\n'.join(edge_lines),
            ' '.join(edge_names),
            elision_str,
            ';\n'.join(above_lines),
            ' -> '.join(order_lines),
         ))

    return s


def dot_graph(vertices, edges,
              directed=False,
              title=None,
              reformat_labels=True,
              vertex_colors=None,
              edge_colors=None,
              edge_labels=None,
              vertex_styles=None,
              edge_styles=None,
              graph_name=None,
              shorten_names=False,
              key_items=None,
              vertex_clusters=None):
    """Generate a Graphviz representation of a list of vertices and edges.

    :param vertices: list of vertex names (optional).
    :param edges:    list of (vertex, vertex) pairs
    :param directed: bool: whether the graph is directed
    :param title: optional title for the graph
    :param reformat_labels: whether to wrap long vertex labels
    :param vertex_colors: if not None, a sequence of colours for the vertices
    :param edge_colors: if not None, colours for the edges
    :param edge_labels: if not None, labels for the edges
    :param vertex_styles: if not None, DOT style strings for vertices
    :param edge_styles: if not None, DOT style strings for edges
    :param graph_name: if not None, name of graph
    :param shorten_names: if True, remove common DN suffixes
    :param key: (is_vertex, style, description) tuples
    :param vertex_clusters: list of subgraph cluster names

    Colour, style, and label lists must be the same length as the
    corresponding list of edges or vertices (or None).

    Colours can be HTML RGB strings ("#FF0000") or common names
    ("red"), or some other formats you don't want to think about.

    If `vertices` is None, only the vertices mentioned in the edges
    are shown, and their appearance can be modified using the
    vertex_colors and vertex_styles arguments. Vertices appearing in
    the edges but not in the `vertices` list will be shown but their
    styles can not be modified.
    """
    out = []
    write = out.append

    if vertices is None:
        vertices = set(x[0] for x in edges) | set(x[1] for x in edges)

    if shorten_names:
        vlist = list(set(x[0] for x in edges) |
                     set(x[1] for x in edges) |
                     set(vertices))
        vmap, elisions = shorten_vertex_names(vlist)
        vertices = [vmap[x] for x in vertices]
        edges = [(vmap[a], vmap[b]) for a, b in edges]

    else:
        elisions = None

    if graph_name is None:
        graph_name = 'A_samba_tool_production'

    if directed:
        graph_type = 'digraph'
        connector = '->'
    else:
        graph_type = 'graph'
        connector = '--'

    write('/* generated by samba */')
    write('%s %s {' % (graph_type, graph_name))
    if title is not None:
        write('label="%s";' % (title,))
    write('fontsize=%s;\n' % (FONT_SIZE))
    write('node[fontname=Helvetica; fontsize=%s];\n' % (FONT_SIZE))

    prev_cluster = None
    cluster_n = 0
    quoted_vertices = []
    for i, v in enumerate(vertices):
        v = quote_graph_label(v, reformat_labels)
        quoted_vertices.append(v)
        attrs = []
        if vertex_clusters and vertex_clusters[i]:
            cluster = vertex_clusters[i]
            if cluster != prev_cluster:
                if prev_cluster is not None:
                    write("}")
                prev_cluster = cluster
                n = quote_graph_label(cluster)
                if cluster:
                    write('subgraph cluster_%d {' % cluster_n)
                    cluster_n += 1
                    write('style = "rounded,dotted";')
                    write('node [style="filled"; fillcolor=white];')
                    write('label = "%s";' % n)

        if vertex_styles and vertex_styles[i]:
            attrs.append(vertex_styles[i])
        if vertex_colors and vertex_colors[i]:
            attrs.append('color="%s"' % quote_graph_label(vertex_colors[i]))
        if attrs:
            write('"%s" [%s];' % (v, ', '.join(attrs)))
        else:
            write('"%s";' % (v,))

    if prev_cluster:
        write("}")

    for i, edge in enumerate(edges):
        a, b = edge
        if a is None:
            a = "Missing source value"
        if b is None:
            b = "Missing destination value"

        a = quote_graph_label(a, reformat_labels)
        b = quote_graph_label(b, reformat_labels)

        attrs = []
        if edge_labels:
            label = quote_graph_label(edge_labels[i])
            attrs.append('label="%s"' % label)
        if edge_colors:
            attrs.append('color="%s"' % quote_graph_label(edge_colors[i]))
        if edge_styles:
            attrs.append(edge_styles[i])  # no quoting
        if attrs:
            write('"%s" %s "%s" [%s];' % (a, connector, b, ', '.join(attrs)))
        else:
            write('"%s" %s "%s";' % (a, connector, b))

    if key_items:
        key = compile_graph_key(key_items, nodes_above=quoted_vertices,
                                elisions=elisions)
        write(key)

    write('}\n')
    return '\n'.join(out)


COLOUR_SETS = {
    'ansi': {
        'alternate rows': (colour.DARK_WHITE, colour.BLACK),
        'disconnected': colour.RED,
        'connected': colour.GREEN,
        'transitive': colour.DARK_YELLOW,
        'header': colour.UNDERLINE,
        'reset': colour.C_NORMAL,
    },
    'ansi-heatmap': {
        'alternate rows': (colour.DARK_WHITE, colour.BLACK),
        'disconnected': colour.REV_RED,
        'connected': colour.REV_GREEN,
        'transitive': colour.REV_DARK_YELLOW,
        'header': colour.UNDERLINE,
        'reset': colour.C_NORMAL,
    },
    'xterm-256color': {
        'alternate rows': (colour.xterm_256_colour(39),
                           colour.xterm_256_colour(45)),
        #'alternate rows': (colour.xterm_256_colour(246),
        #                   colour.xterm_256_colour(247)),
        'disconnected': colour.xterm_256_colour(124, bg=True),
        'connected': colour.xterm_256_colour(112),
        'transitive': colour.xterm_256_colour(214),
        'transitive scale': (colour.xterm_256_colour(190),
                             colour.xterm_256_colour(184),
                             colour.xterm_256_colour(220),
                             colour.xterm_256_colour(214),
                             colour.xterm_256_colour(208),
        ),
        'header': colour.UNDERLINE,
        'reset': colour.C_NORMAL,
    },
    'xterm-256color-heatmap': {
        'alternate rows': (colour.xterm_256_colour(171),
                           colour.xterm_256_colour(207)),
        #'alternate rows': (colour.xterm_256_colour(246),
        #                    colour.xterm_256_colour(247)),
        'disconnected': colour.xterm_256_colour(124, bg=True),
        'connected': colour.xterm_256_colour(112, bg=True),
        'transitive': colour.xterm_256_colour(214, bg=True),
        'transitive scale': (colour.xterm_256_colour(190, bg=True),
                             colour.xterm_256_colour(184, bg=True),
                             colour.xterm_256_colour(220, bg=True),
                             colour.xterm_256_colour(214, bg=True),
                             colour.xterm_256_colour(208, bg=True),
        ),
        'header': colour.UNDERLINE,
        'reset': colour.C_NORMAL,
    },
    None: {
        'alternate rows': ('',),
        'disconnected': '',
        'connected': '',
        'transitive': '',
        'header': '',
        'reset': '',
    }
}

CHARSETS = {
    'utf8': {
        'vertical': '│',
        'horizontal': '─',
        'corner': '╭',
        #'diagonal': '╲',
        'diagonal': '·',
        #'missing': '🕱',
        'missing': '-',
        'right_arrow': '←',
    },
    'ascii': {
        'vertical': '|',
        'horizontal': '-',
        'corner': ',',
        'diagonal': '0',
        'missing': '-',
        'right_arrow': '<-',
    }
}


def find_transitive_distance(vertices, edges):
    all_vertices = (set(vertices) |
                    set(e[0] for e in edges) |
                    set(e[1] for e in edges))

    if all_vertices != set(vertices):
        print("there are unknown vertices: %s" %
              (all_vertices - set(vertices)),
              file=sys.stderr)

    # with n vertices, we are always less than n hops away from
    # anywhere else.
    inf = len(all_vertices)
    distances = {}
    for v in all_vertices:
        distances[v] = {v: 0}

    for src, dest in edges:
        distances[src][dest] = distances[src].get(dest, 1)

    # This algorithm (and implementation) seems very suboptimal.
    # potentially O(n^4), though n is smallish.
    for i in range(inf):
        changed = False
        new_distances = {}
        for v, d in distances.items():
            new_d = d.copy()
            new_distances[v] = new_d
            for dest, cost in d.items():
                for leaf, cost2 in distances[dest].items():
                    new_cost = cost + cost2
                    old_cost = d.get(leaf, inf)
                    if new_cost < old_cost:
                        new_d[leaf] = new_cost
                        changed = True

        distances = new_distances
        if not changed:
            break

    # filter out unwanted vertices and infinite links
    answer = {}
    for v in vertices:
        answer[v] = {}
        for v2 in vertices:
            a = distances[v].get(v2, inf)
            if a < inf:
                answer[v][v2] = a

    return answer


def get_transitive_colourer(colours, n_vertices):
    if 'transitive scale' in colours:
        scale = colours['transitive scale']
        m = len(scale)
        n = 1 + int(n_vertices ** 0.5)

        def f(link):
            if not isinstance(link, int):
                return ''
            return scale[min(link * m // n, m - 1)]

    else:
        def f(link):
            return colours['transitive']

    return f


def distance_matrix(vertices, edges,
                    utf8=False,
                    colour=None,
                    shorten_names=False,
                    generate_key=False,
                    grouping_function=None,
                    row_comments=None):
    lines = []
    write = lines.append

    charset = CHARSETS['utf8' if utf8 else 'ascii']
    vertical = charset['vertical']
    horizontal = charset['horizontal']
    corner = charset['corner']
    diagonal = charset['diagonal']
    missing = charset['missing']
    right_arrow = charset['right_arrow']

    colours = COLOUR_SETS[colour]

    colour_cycle = cycle(colours.get('alternate rows', ('',)))

    if vertices is None:
        vertices = sorted(set(x[0] for x in edges) | set(x[1] for x in edges))

    if grouping_function is not None:
        # we sort and colour according to the grouping function
        # which can be used to e.g. alternate colours by site.
        vertices = sorted(vertices, key=grouping_function)
        colour_list = []
        for k, v in groupby(vertices, key=grouping_function):
            c = next(colour_cycle)
            colour_list.extend(c for x in v)
    else:
        colour_list = [next(colour_cycle) for v in vertices]

    if shorten_names:
        vlist = list(set(x[0] for x in edges) |
                     set(x[1] for x in edges) |
                     set(vertices))
        vmap, replacements = shorten_vertex_names(vlist, '+',
                                                  aggressive=True)
        vertices = [vmap[x] for x in vertices]
        edges = [(vmap[a], vmap[b]) for a, b in edges]


    vlen = max(6, max(len(v) for v in vertices))

    # first, the key for the columns
    c_header = colours.get('header', '')
    c_disconn = colours.get('disconnected', '')
    c_conn = colours.get('connected', '')
    c_reset = colours.get('reset', '')

    colour_transitive = get_transitive_colourer(colours, len(vertices))

    vspace = ' ' * vlen
    verticals = ''
    write("%*s %s  %sdestination%s" % (vlen, '',
                                       ' ' * len(vertices),
                                       c_header,
                                       c_reset))
    for i, v in enumerate(vertices):
        j = len(vertices) - i
        c = colour_list[i]
        if j == 1:
            start = '%s%ssource%s' % (vspace[:-6], c_header, c_reset)
        else:
            start = vspace
        write('%s %s%s%s%s%s %s%s' % (start,
                                      verticals,
                                      c_reset,
                                      c,
                                      corner,
                                      horizontal * j,
                                      v,
                                      c_reset
        ))
        verticals += c + vertical

    connections = find_transitive_distance(vertices, edges)

    for i, v in enumerate(vertices):
        c = colour_list[i]
        links = connections[v]
        row = []
        for v2 in vertices:
            link = links.get(v2)
            if link is None:
                row.append('%s%s' % (c_disconn, missing))
                continue
            if link == 0:
                row.append('%s%s%s%s' % (c_reset, c, diagonal, c_reset))
            elif link == 1:
                row.append('%s1%s' % (c_conn, c_reset))
            else:
                ct = colour_transitive(link)
                if link > 9:
                    link = '>'
                row.append('%s%s%s' % (ct, link, c_reset))

        if row_comments is not None and row_comments[i]:
            row.append('%s %s %s' % (c_reset, right_arrow, row_comments[i]))

        write('%s%*s%s %s%s' % (c, vlen, v, c_reset,
                                ''.join(row), c_reset))

    example_c = next(colour_cycle)
    if shorten_names:
        write('')
        for substitute, original in reversed(replacements):
            write("'%s%s%s' stands for '%s%s%s'" % (example_c,
                                                    substitute,
                                                    c_reset,
                                                    example_c,
                                                    original,
                                                    c_reset))
    if generate_key:
        write('')
        write("Data can get from %ssource%s to %sdestination%s in the "
              "indicated number of steps." % (c_header, c_reset,
                                              c_header, c_reset))
        write("%s%s%s means zero steps (it is the same DC)" %
              (example_c, diagonal, c_reset))
        write("%s1%s means a direct link" % (c_conn, c_reset))
        write("%s2%s means a transitive link involving two steps "
              "(i.e. one intermediate DC)" %
              (colour_transitive(2), c_reset))
        write("%s%s%s means there is no connection, even through other DCs" %
              (c_disconn, missing, c_reset))

    return '\n'.join(lines)


def pad_char(char, digits, padding=' '):
    if digits == 1:
        padding = ''
    return ' ' * (digits - 1) + char + padding


def transpose_dict_matrix(m):
    m2 = {}
    for k1, row in m.items():
        for k2, dist in row.items():
            m2.setdefault(k2, {})[k1] = dist
    return m2

def full_matrix(rows,
                utf8=False,
                colour=None,
                shorten_names=False,
                generate_key=False,
                grouping_function=None,
                row_comments=None,
                colour_scale=None,
                digits=1,
                ylabel='source',
                xlabel='destination',
                transpose=True):
    lines = []
    write = lines.append

    if transpose:
        rows = transpose_dict_matrix(rows)

    use_padding = digits > 1

    charset = CHARSETS['utf8' if utf8 else 'ascii']
    vertical = pad_char(charset['vertical'], digits)
    horizontal = charset['horizontal'] * (digits + use_padding)
    corner = pad_char(charset['corner'], digits,
                      charset['horizontal'])
    diagonal = pad_char(charset['diagonal'], digits)
    missing = pad_char(charset['missing'], digits)
    toobig = pad_char('>', digits)
    right_arrow = charset['right_arrow']
    empty = pad_char(' ', digits)

    colours = COLOUR_SETS[colour]

    colour_cycle = cycle(colours.get('alternate rows', ('',)))
    vertices = list(rows.keys())
    if grouping_function is not None:
        # we sort and colour according to the grouping function
        # which can be used to e.g. alternate colours by site.
        vertices.sort(key=grouping_function)
        colour_list = []
        for k, v in groupby(vertices, key=grouping_function):
            c = next(colour_cycle)
            colour_list.extend(c for x in v)
    else:
        colour_list = [next(colour_cycle) for v in vertices]

    if shorten_names:
        vmap, replacements = shorten_vertex_names(vertices, '+',
                                                  aggressive=True)
        rows2 = {}
        for vert, r in rows.items():
            rows2[vmap[vert]] = dict((vmap[k], v) for k, v in r.items())

        rows = rows2
        vertices = list(rows.keys())

    vlen = max(6, len(xlabel), max(len(v) for v in vertices))

    # first, the key for the columns
    c_header = colours.get('header', '')
    c_disconn = colours.get('disconnected', '')
    c_conn = colours.get('connected', '')
    c_reset = colours.get('reset', '')

    if colour_scale is None:
        colour_scale = len(rows)
    colour_transitive = get_transitive_colourer(colours, colour_scale)

    vspace = ' ' * vlen
    verticals = ''
    write("%s %s %s%s%s" % (vspace,
                            empty * (len(rows) + 1),
                            c_header,
                            xlabel,
                            c_reset))
    for i, v in enumerate(vertices):
        j = len(rows) - i
        c = colour_list[i]
        if j == 1:
            start = '%s%s%s%s' % (vspace[:-len(ylabel)],
                                  c_header,
                                  ylabel,
                                  c_reset)
        else:
            start = vspace
        write('%s %s%s%s%s%s %s%s' % (start,
                                      verticals,
                                      c_reset,
                                      c,
                                      corner,
                                      horizontal * j,
                                      v,
                                      c_reset
        ))
        verticals += '%s%s' % (c, vertical)

    end_cell = '%s%s' % (' ' * use_padding, c_reset)
    overflow = False
    for i, v in enumerate(vertices):
        links = rows[v]
        c = colour_list[i]
        row = []
        for v2 in vertices:
            if v2 not in links:
                row.append('%s%s%s' % (c_disconn, missing, c_reset))
            elif v == v2:
                row.append('%s%s%s%s' % (c_reset, c, diagonal, c_reset))
            else:
                link = links[v2]
                if link >= 10 ** digits:
                    ct = colour_transitive(link)
                    row.append('%s%s%s' % (ct, toobig, c_reset))
                    overflow = True
                    continue
                if link == 0:
                    ct = c_conn
                else:
                    ct = colour_transitive(link)
                row.append('%s%*s%s' % (ct, digits, link, end_cell))

        if row_comments is not None and row_comments[i]:
            row.append('%s %s %s' % (c_reset, right_arrow, row_comments[i]))

        write('%s%*s%s %s%s' % (c, vlen, v, c_reset,
                                ''.join(row), c_reset))

    if overflow or shorten_names:
        write('')

    if overflow:
            write("'%s%s%s' means greater than %d " %
                  (colour_transitive(10 ** digits),
                   toobig,
                   c_reset,
                   10 ** digits - 1))

    if shorten_names:
        example_c = next(colour_cycle)
        for substitute, original in reversed(replacements):
            write("'%s%s%s' stands for '%s%s%s'" % (example_c,
                                                    substitute,
                                                    c_reset,
                                                    example_c,
                                                    original,
                                                    c_reset))

    return '\n'.join(lines)