path: root/lib.py

#!/usr/bin/env python
# vim: set fileencoding=UTF-8:
# Copyright 2013 Red Hat, Inc.
# Author: Jan Pokorný <jpokorny at redhat dot com>
# Distributed under GPLv2+;  generated content under CC-BY-SA 3.0
# (to view a copy, visit http://creativecommons.org/licenses/by-sa/3.0/)
"""Elements library incl. visual aspects"""

from sys import path
from os.path import expanduser, extsep
from pydot import Dot, Edge, Node, Subgraph

__all__ = []

# modify/mutate BLACKLIST in order to filter some elements out
# (not a beautiful, rather ad-hoc design, I know)
BLACKLIST = []


#
# Customization via meta level
#

# obfuscation contest in recursively yielding base classes up to given set
bases = lambda x: (lambda f_, *xs_: f_(f_, *xs_)) \
                  (lambda f, x=[], *xs:
                   [x] + f(f, *x.__bases__) + f(f, *xs)
                   if x not in ([], LibSubgraph, LibNode, LibEdge)
                   else [],
                   x)
bl_test = lambda ix: not set(bases(ix.__class__)).intersection(BLACKLIST)


def bl_map_nodes_edges(nodes, edges, prev_nnames=None):
    nnames = prev_nnames if prev_nnames is not None else set()
    for n in nodes:
        n.get_attributes()['style'] = 'invis'
        nnames.add(n.get_name())
    for e in edges:
        if not bl_test(e) \
                or e.get_source() in nnames or e.get_destination() in nnames:
            e.get_attributes()['style'] = 'invis'


class LibMeta(type):
    def __init__(cls, name, bases, attrs):
        __all__.append(name)
        super(LibMeta, cls).__init__(name, bases, attrs)
        old_init = cls.__init__
        def new_init(self, *args, **kwargs):
            subgraphs = kwargs.pop('_subgraphs', ())
            nodes = kwargs.pop('_nodes', ())
            edges = kwargs.pop('_edges', ())

            if not hasattr(self, 'nnames'):
                self.nnames = set()
            nnames = self.nnames
            map(lambda x: nnames.update(getattr(x, 'nnames')),
                tuple(subgraphs) + tuple(nodes) + tuple(edges))

            old_init(self, *args, **kwargs)

            bl_map_nodes_edges(filter(lambda x: not bl_test(x), nodes),
                               edges, nnames)
            for n in nodes:
                self.add_node(n)
            for e in edges:
                self.add_edge(e)

            for s in filter(bl_test, subgraphs):
                s and self.add_subgraph(s)
        cls.__init__ = new_init


class LibDot(Dot):
    __metaclass__ = LibMeta
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('graph_type', 'digraph')
        super(LibDot, self).__init__(*args, **kwargs)
        #self.set_suppress_disconnected(True)


class LibSubgraph(Subgraph):
    __metaclass__ = LibMeta
    def __init__(self, *args, **kwargs):
        super(LibSubgraph, self).__init__(*args, **kwargs)
        #self.set_suppress_disconnected(True)


class LibNode(Node):
    __metaclass__ = LibMeta


class LibEdge(Edge):
    __metaclass__ = LibMeta


#
# Subgraph inheritance
#


class SubgraphImportant(LibSubgraph):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('fillcolor', '#eaeaea')
        kwargs.setdefault('style', 'bold, filled, rounded')
        super(SubgraphImportant, self).__init__(*args, **kwargs)


class SubgraphStandard(LibSubgraph):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('fillcolor', '#f5f5f5')
        kwargs.setdefault('style', 'filled, rounded')
        super(SubgraphStandard, self).__init__(*args, **kwargs)


class SubgraphInvisible(LibSubgraph):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('style', 'invis')
        super(SubgraphInvisible, self).__init__(*args, **kwargs)



#
# Node inheritance
#


class NodeInvisible(LibNode):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('style', 'invis')
        super(LibNode, self).__init__(*args, **kwargs)


class Program(LibNode):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('shape', 'box')
        kwargs.setdefault('style', 'filled')
        super(Program, self).__init__(*args, **kwargs)


class Agent(Program):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('fillcolor', 'lavenderblush')
        kwargs.setdefault('style', 'filled')
        super(Agent, self).__init__(*args, **kwargs)


class Library(Program):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('fillcolor', 'lavenderblush')
        kwargs.setdefault('style', 'filled')
        super(Library, self).__init__(*args, **kwargs)


class Executable(Program):
    pass


class Daemon(Executable):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('fillcolor', 'cornsilk')
        kwargs.setdefault('style', 'filled')
        super(Daemon, self).__init__(*args, **kwargs)


class Artefact(LibNode):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('shape', 'box3d')
        kwargs.setdefault('fillcolor', 'wheat')
        kwargs.setdefault('style', 'filled')
        super(Artefact, self).__init__(*args, **kwargs)


class Device(LibNode):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('shape', 'box3d')
        kwargs.setdefault('fillcolor', 'wheat')
        kwargs.setdefault('style', 'filled')
        super(Device, self).__init__(*args, **kwargs)


class StorageDevice(Device):
    def __init__(self, *args, **kwargs):
        super(StorageDevice, self).__init__(*args, **kwargs)


class FenceDevice(Device):
    def __init__(self, *args, **kwargs):
        super(FenceDevice, self).__init__(*args, **kwargs)


#
# Edge inheritance
#

class EdgeInvisible(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('style', 'invis')
        kwargs.setdefault('arrowhead', 'none')
        super(EdgeInvisible, self).__init__(*args, **kwargs)


class DBUS(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('style', 'dotted')
        super(DBUS, self).__init__(*args, **kwargs)


class OddjobExec(DBUS):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('label', 'oddjob\nexec')
        super(OddjobExec, self).__init__(*args, **kwargs)


class HTTPS(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('style', 'bold')
        super(HTTPS, self).__init__(*args, **kwargs)


class LuciHTTPS(HTTPS):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('label', 'port 8084')
        super(LuciHTTPS, self).__init__(*args, **kwargs)


class ApplicationSpecificProtocol(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('style', 'dashed')
        super(ApplicationSpecificProtocol, self).__init__(*args, **kwargs)


class RICCIRPC(ApplicationSpecificProtocol):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('label', 'port 11111')
        super(RICCIRPC, self).__init__(*args, **kwargs)


class SNMP(ApplicationSpecificProtocol):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('label', 'port 161')
        super(SNMP, self).__init__(*args, **kwargs)


class CIM(ApplicationSpecificProtocol):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('label', 'port 898[89]')
        super(CIM, self).__init__(*args, **kwargs)


class Consume(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('color', 'darkgreen')
        kwargs.setdefault('fontcolor', 'darkgreen')
        super(Consume, self).__init__(*args, **kwargs)


class ConsumeReversed(Consume):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('dir', 'back')
        super(ConsumeReversed, self).__init__(*reversed(args), **kwargs)


class Produce(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('color', 'orangered')
        kwargs.setdefault('fontcolor', 'orangered')
        super(Produce, self).__init__(*args, **kwargs)


class Delegate(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('color', 'navy')
        kwargs.setdefault('fontcolor', 'navy')
        #kwargs.setdefault('dir', 'both')
        super(Delegate, self).__init__(*args, **kwargs)


class FencedBy(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('color', 'red')
        kwargs.setdefault('dir', 'back')
        super(FencedBy, self).__init__(*args, **kwargs)


# peer-to-peer (1:1 hence constraint='false')
class Exchange(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('color', 'saddlebrown')
        kwargs.setdefault('fontcolor', 'saddlebrown')
        kwargs.setdefault('dir', 'both')
        kwargs.setdefault('constraint', 'false')
        super(Exchange, self).__init__(*args, **kwargs)


# data bus
class Databus(Exchange):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('color', 'saddlebrown')
        kwargs.setdefault('fontcolor', 'saddlebrown')
        super(Databus, self).__init__(*args, **kwargs)


# mixing

class DelegateOddjobExec(Delegate, OddjobExec):
    pass


class DelegateRICCIRPC(Delegate, RICCIRPC):
    pass


class DelegateLuciHTTPS(Delegate, LuciHTTPS):
    pass


class DelegateSNMP(Delegate, SNMP):
    pass


class DelegateCIM(Delegate, CIM):
    pass


#
# main-helpers
#

def export(fnc):
    __all__.append(fnc.__name__)
    return fnc


@export
def gen_graph(graph, blacklist=(), **kwargs):
    # can eventually restore the state
    BLACKLIST[:] = blacklist
    return graph()


@export
def xdot_graph(*args, **kwargs):
    import gtk
    import gtk.gdk
    try:
        path[:] = [expanduser('~/wrkspc/goss-medium/jrfonseca.xdot')] + path[:]
        import xdot
    except ImportError:
        print 'missing xdot; use "pip install xdot" or equivalent'
        raise

    class LibDotWindow(xdot.DotWindow):
        # heavily inspired from http://zetcode.com/gui/pygtk/menus/
        def on_toggle_item(self, widget):
            label = widget.get_label()
            bl = self._kwargs.setdefault('blacklist', [])
            change = False
            cls = globals()[label]
            if widget.active and cls in bl:
                bl.remove(cls)
                bs = bases(cls)
                bs.remove(cls)
                for item in filter(lambda x: isinstance(x, gtk.CheckMenuItem)
                                   and globals()[x.get_label()] in bs,
                                   widget.get_parent().get_children()):
                    if not item.get_active():
                        item.set_active(True)
                change = True
            elif not widget.active and cls not in bl:
                bl.append(cls)
                change = True
            if change and self._can_change:
                self.set_dotcode(
                    gen_graph(*self._args, **self._kwargs).to_string())

        def on_clean_all(self, widget):
            self._can_change = False
            for item in filter(lambda x: isinstance(x, gtk.CheckMenuItem),
                               widget.get_parent().get_children()):
                item.set_active(False)
            self.set_dotcode(
                gen_graph(*self._args, **self._kwargs).to_string())
            self._can_change = True

        def on_set_all(self, widget):
            self._can_change = False
            for item in filter(lambda x: isinstance(x, gtk.CheckMenuItem),
                               widget.get_parent().get_children()):
                item.set_active(True)
            self.set_dotcode(
                gen_graph(*self._args, **self._kwargs).to_string())
            self._can_change = True

        def __init__(self, *args, **kwargs):
            self._args = args
            self._kwargs = kwargs
            self._can_change = True
            super(LibDotWindow, self).__init__()
            dotcode = gen_graph(*args, **kwargs).to_string()
            title = dotcode[dotcode.find('digraph', 0, 512) + len('digraph')
                            :
                            dotcode.find('{', 0, 512)].strip(' \'"')
            self.base_title = title
            self.set_dotcode(dotcode)

            mb = gtk.MenuBar()

            filemenu = gtk.Menu()
            filem = gtk.MenuItem("File")
            filem.set_submenu(filemenu)

            edgemenu = gtk.Menu()
            edge = gtk.MenuItem("Edge")
            edge.set_submenu(edgemenu)
            edgemenuitem = gtk.MenuItem("Clean all")
            edgemenuitem.connect("activate", self.on_clean_all)
            edgemenu.append(edgemenuitem)
            edgemenuitem = gtk.MenuItem("Set all")
            edgemenuitem.connect("activate", self.on_set_all)
            edgemenu.append(edgemenuitem)
            edgemenu.append(gtk.SeparatorMenuItem())

            nodemenu = gtk.Menu()
            node = gtk.MenuItem("Node")
            node.set_submenu(nodemenu)
            nodemenuitem = gtk.MenuItem("Clean all")
            nodemenuitem.connect("activate", self.on_clean_all)
            nodemenu.append(nodemenuitem)
            nodemenuitem = gtk.MenuItem("Set all")
            nodemenuitem.connect("activate", self.on_set_all)
            nodemenu.append(nodemenuitem)
            nodemenu.append(gtk.SeparatorMenuItem())

            for i in __all__:
                try:
                    i = globals()[i]
                    if issubclass(i, LibEdge):
                        if i in (LibEdge, EdgeInvisible):
                            continue
                        imenuitem = gtk.CheckMenuItem("%s" % i.__name__)
                        imenuitem.set_active(True)
                        imenuitem.connect("activate", self.on_toggle_item)
                        edgemenu.append(imenuitem)
                    elif issubclass(i, LibNode):
                        if i in (LibNode, NodeInvisible):
                            continue
                        imenuitem = gtk.CheckMenuItem("%s" % i.__name__)
                        imenuitem.set_active(True)
                        imenuitem.connect("activate", self.on_toggle_item)
                        nodemenu.append(imenuitem)
                except:
                    pass

            exit = gtk.MenuItem("Exit")
            exit.connect("activate", gtk.main_quit)
            filemenu.append(exit)

            mb.append(filem)
            mb.append(edge)
            mb.append(node)

            sb = gtk.Statusbar()
            sb.push(1, "Ready")

            vbox = self.get_children()[0]
            vbox.pack_start(mb, False, False, 0)
            vbox.pack_start(sb, False, False, 0)
            vbox.reorder_child(mb, 0)

            self.connect("destroy", gtk.main_quit)
            self.maximize()
            self.show_all()

    window = LibDotWindow(*args, **kwargs)
    window.connect('destroy', gtk.main_quit)
    gtk.main()


@export
def main(graph, argv, *args, **kws):
    # -x deprecated, supported for backward compatibility
    x, argv = (1, argv[1:]) if len(argv) <= 1 or argv[1] == '-x' else (0, argv)
    output = kws.pop('output', 'sinenomine')
    if x:
        xdot_graph(graph, **kws)
    else:
        ext = argv[1] if len(argv) > 1 else 'pdf'
        fmt = {'dot': 'raw'}.get(ext, ext)
        output += extsep + ext
        gen_graph(graph, **kws).write(output, format=fmt, prog=['dot'].extend(args))