source: fedd/fedd_image.py @ 9252414

#!/usr/local/bin/python

import sys
import os
import tempfile
import subprocess
import xml.parsers.expat

from federation import topdl
from federation.remote_service import service_caller
from federation.client_lib import client_opts, exit_with_fault, RPCException, \
        wrangle_standard_options, do_rpc, get_experiment_names, save_certfile


class image_opts(client_opts):
    def __init__(self):
        client_opts.__init__(self)
        self.add_option("--experiment_cert", dest="exp_certfile",
                type="string", help="experiment certificate file")
        self.add_option("--experiment_name", dest="exp_name",
                type="string", help="human readable experiment name")
        self.add_option("--output", dest="outfile", type="string",
                help="output image file")
        self.add_option("--format", dest="format", type="choice", 
                choices=("jpg", "png", "dot", "svg"),
                help="Output file format override")
        self.add_option("--program", dest="neato", default=None,
                type="string",
                help="Program compatible with dot (from graphviz) used to " + \
                        "render image")
        self.add_option("--labels", dest='labels', action='store_true',
                default=True, help='Label nodes and edges')
        self.add_option("--no_labels", dest='labels',
                default=True, action='store_false',
                help='Label nodes and edges')
        self.add_option('--pixels', dest="pixels", default=None,
                type="int",
                help="Size of output in pixels (diagrams are square")
        self.add_option("--file", dest="file", 
                help="experiment description file")
        self.add_option("--group", dest="group", action="append", default=[], 
                help="Group nodes by this attribute")

def make_subgraph(topelems, attrs):
    """
    Take a list of topdl elements (Computers and Substrates ) and partition
    them into groups based on the value of the given attribute.  All computers
    with the same value for the attribute are in a subgraph. A substrate is in
    the subgraph if all its interfaces are, and in the "default" subgraph if
    not.
    """

    def add_to_map(m, a, v):
        """
        Add this value to the list in m[a] if that list exists, or create it.
        """
        if a in m: m[a].append(v)
        else: m[a] = [ v]

    if not attrs: return { }

    sg = { }
    attr, rest = attrs[0], attrs[1:]
    for e in topelems:
        if isinstance(e, topdl.Computer):
            add_to_map(sg, e.get_attribute(attr) or 'default', e)
        elif isinstance(e, topdl.Substrate):
            # A little hairy.  Make a set of the values of the given attribute
            # for all the elements in the substrate's interfaces.  If that set
            # has one element, put the substrate in that subgraph.  Otherwise
            # put it in default.
            atts = set([i.element.get_attribute(attr) or 'default' \
                    for i in e.interfaces])

            if len(atts) == 1: add_to_map(sg, atts.pop(), e)
            else: add_to_map(sg, 'default', e)

    if rest:
        for k in sg.keys():
            sg[k] = make_subgraph(sg[k], rest)

    return sg

def output_subgraphs(sg, dotfile):

    # Construct subgraphs featurning nodes from those graphs
    for sn, nodes in sg.items():
        # Output the computers and lans in subgraphs, if necessary.  Subgraphs
        # treated as clusters are designated by "cluster" starting the name.
        # (The default subgraph is not treated as a cluster.
        if sn != 'default': 
            print >>dotfile, 'subgraph "cluster-%s" {' % sn
            print >>dotfile, 'label="%s"' % sn
        else: 
            print >>dotfile, 'subgraph {'
        print >>dotfile, 'color="#08c0f8"'
        print >>dotfile, 'clusterrank="local"'

        if isinstance(nodes, dict): output_subgraphs(nodes, dotfile)
        elif isinstance(nodes, list):
            # For each node in the subgraph, output its representation
            for i, n in enumerate(nodes):
                if isinstance(n, topdl.Computer):
                    if n.name:
                        print >>dotfile, \
                                '\t"%s" [shape=box,style=filled,\\' % n.name
                    else:
                        print >>dotfile, \
                                '\t"unnamed_node%d" [shape=box,style=filled,\\'\
                                % i
                    print >>dotfile, \
                            '\t\tcolor=black,fillcolor="#60f8c0",regular=1]'
                elif isinstance(n, topdl.Substrate):
                    print >>dotfile, '\t"%s" [shape=ellipse, style=filled,\\' \
                            % n.name
                    print >>dotfile,\
                            '\t\tcolor=black,fillcolor="#80c0f8",regular=1]'
        print >>dotfile, '}'

def gen_dot_topo(t, labels, dotfile, sg_attr):
    """
    Write a dot description of the topology in t, a topdl.Topology to the open
    file dotfile.  the gen_image function below has already put some of that
    file together. This function only generates the nodes and connections.  If
    labels is true, label the nodes.
    """


    # The routine draws a circle for substrates with more than 2 interfaces
    # connected to it, so this makes lists of each of those.
    lans = [ s for s in t.substrates if len(s.interfaces) != 2]
    links = [ s for s in t.substrates if len(s.interfaces) == 2]

    # Break into subgraphs by the attribute of one's given
    if sg_attr: sg = make_subgraph(t.elements + lans, sg_attr)
    else: sg = { 'default': t.elements + lans }


    output_subgraphs(sg, dotfile)

    # Pull the edges out ot the lans and links
    for l in lans:
        for i in l.interfaces:
            ip = i.get_attribute('ip4_address')
            if labels and ip:
                print >>dotfile, '\t"%s" -- "%s" [headlabel="%s"]' % \
                        (l.name, i.element.name, ip)
            else:
                print >>dotfile, '\t"%s" -- "%s"' % \
                        (l.name, i.element.name)

    for l in links:
        s, d = l.interfaces[0:2] 
        sip = s.get_attribute('ip4_address')
        dip = d.get_attribute('ip4_address')
        if labels and sip and dip:
            print >>dotfile, \
                    ('\t"%s" -- "%s" [label="%s",taillabel="%s",' + \
                    'headlabel="%s"]') % \
                    (s.element.name, d.element.name, l.name,
                        sip, dip)
        else:
            print >>dotfile, '\t"%s" -- "%s" ' % \
                    (s.element.name, d.element.name)

def gen_image(top, file, fmt, neato, labels, pix=None, sg_attr=None):
    """
    Create a dot formatted input file from the topdl.Topology in top and run
    dot on it, storing the output in a file named file.  The image format it in
    fmt, and it has to be one of the formats understood by dot or neato.
    Labels is true if the graph should be noted with node names and addresses.
    Pix sets the sixe of the output in pixels (square images).  Neato is a
    user-specified graph drawing programs that takes neato/dot compaitble
    inputs.  If no neato is specified, the code looks for a suitable
    executable.  Return the command's return code.
    """

    # Open up a temporary file for dot to turn into a visualization.  This
    # may raise an exception
    df, dotname = tempfile.mkstemp(prefix='fedd_client', suffix=".dot")
    dotfile = os.fdopen(df, 'w')

    # Look for neato if the executable is unspecified.
    if not neato:
        # If a subgraph attribute has been specified, prefer dot to neato.  dot
        # is more "natural" looking graphs, but neato segregates better.
        if sg_attr: 
            search_list = ['/usr/bin/dot', '/usr/local/bin/dot',
                    '/usr/bin/neato', '/usr/local/bin/neato']
        else:
            search_list = ['/usr/bin/neato', '/usr/local/bin/neato', 
                '/usr/bin/dot', '/usr/local/bin/dot']

        for f in search_list:
            if os.access(f, os.X_OK):
                neato = f
                break
        else:
            raise RuntimeError("Cannot find graph rendering program")

    # Some arguments bases on size, outputformat, destination, etc.
    cmd = [neato, '-Gsplines=true']
    if fmt != 'dot': cmd.append('-T%s' % fmt)
    if file:
        cmd.append('-o')
        cmd.append(file)
    cmd.append(dotname)

    nodes = len(top.elements)
    if nodes < 10: size = 5
    elif nodes < 50: size = 10
    else: size = 18

    if pix:
        dpi = pix / size
    else:
        dpi = None

    # Begin composing the graph description
    print >>dotfile, "graph G {"
    if dpi:
        print >>dotfile, \
            '\tgraph [size="%i,%i", dpi="%i", ratio=fill, compound="true"];' \
            % (size, size, dpi)
    else:
        print >>dotfile, \
            '\tgraph [size="%i,%i", ratio=fill, compound="true"];' \
            % (size, size)

    if labels:
        print >>dotfile, '\tnode [fontname=arial,fontsize=9,label="\N"];'
        print >>dotfile, '\tedge [fontname=arial,fontsize=9];\n'
    else:
        print >>dotfile, '\tnode [label=""];'

    gen_dot_topo(top, labels, dotfile, sg_attr=sg_attr)
    print >>dotfile, "}"
    dotfile.close()

    # Redirect the drawing program stderr, and draw
    dev_null = open("/dev/null", "w")
    rv = subprocess.call(cmd, stderr=dev_null)
    os.remove(dotname)
    dev_null.close()

    return rv

def extract_topo_from_message(msg, serialize):
    """
    Extract the topdl.Topology from either an Ns2Topdl or Info response.  In
    the unlikely event that the message is misformatted, a RuntimeError is
    raised.  If the overall script is only serializing messages, this routine
    probably returns None (unless do_rpc is modified to return a valid response
    when serializing only).
    """
    if 'experimentdescription' in msg:
        if 'topdldescription' in msg['experimentdescription']:
            exp = msg['experimentdescription']['topdldescription']
            return topdl.Topology(**exp)
        else:
            raise RuntimeError("Bad response: could not translate")
    elif serialize:
        return None
    else:
        raise RuntimeError("Bad response. %s" % e.message)

def get_experiment_topo(opts, cert, url):
    """
    Get the topology of an existing experiment from a running fedd.  That is,
    pull the experimentdescription from an Info operation. Extract the
    experiment identifier from opts and raise a RuntimeError if insufficient or
    ambiguous information is present there.  The Topology is returned.
    """

    # Figure out what experiment to retrieve
    if opts.exp_name and opts.exp_certfile:
        raise RuntimeError("Only one of --experiment_cert and " + \
                "--experiment_name permitted")
    elif opts.exp_certfile:
        exp_id = { 'fedid': fedid(file=opts.exp_certfile) }
    elif opts.exp_name:
        exp_id = { 'localname' : opts.exp_name }
    else:
        raise RuntimeError("specify one of --experiment_cert, " + \
                "--experiment_name, or --file")

    # Get the topology and return it
    req = { 'experiment': exp_id }
    resp_dict = do_rpc(req,
            url, opts.transport, cert, opts.trusted, 
            serialize_only=opts.serialize_only,
            tracefile=opts.tracefile,
            caller=service_caller('Info'), responseBody="InfoResponseBody")

    return extract_topo_from_message(resp_dict, opts.serialize_only)

def get_file_topo(opts, cert, url):
    """
    Parse a topology file or convert an ns2 file into a topdl.Topology.  XML
    encodings of the topology can be converted completely locally, but ns2
    files must be converted by a fedd.  The topology is returned or a
    RuntimeError, RPCError, or EnvironmentError (if the file is unreadable) is
    raised.
    """
    # Try the easy case first.  Note that if the file is unreadable rather than
    # unparsable, this attempt will raise an EnvironmentError.
    try:
        return topdl.topology_from_xml(filename=opts.file, top='experiment')
    except xml.parsers.expat.ExpatError, e:
        # Print an error message if debugging
        if opts.debug > 1: print >>sys.stderr, "XML parse failed: %s" %e

    # If we get here the parser above failed and we assume that the file is
    # ns2/tcl.  Call out to the converter.  The exceptions above can come out
    # of here, though EnvironmentErrors are unlikely. 
    if not cert:
        raise RuntimeError("Couldn't find an identity certificate, " + \
                "RPC (Ns2Topdl) needed.")

    contents = "".join([l for l in open(opts.file, "r")])
    msg = { 'description': { 'ns2description': contents }, }
    if opts.debug > 1: print >>sys.stderr, msg
    resp_dict = do_rpc(msg, 
            url, opts.transport, cert, opts.trusted, 
            serialize_only=opts.serialize_only,
            tracefile=opts.tracefile,
            caller=service_caller('Ns2Topdl'),
            responseBody="Ns2TopdlResponseBody")

    return extract_topo_from_message(resp_dict, opts.serialize_only)

parser = image_opts()
(opts, args) = parser.parse_args()

# Note that if we're converting a local topdl file, the absence of a cert may
# not be a big deal.  It's checked if needed below
try:
    cert, fid, url = wrangle_standard_options(opts)
except RuntimeError, e:
    cert = None
    print >>sys.stderr, "Warning: %s e"


# Figure out where output is going and in what format.  If format is given and
# understood, that's the output format.  Otherwise we try the last three
# letters of the file as a suffix.  If all that fails, we fall back to
# dot/neato format.
if opts.format and opts.outfile:
    fmt = opts.format
    file = opts.outfile
elif not opts.format and opts.outfile:
    fmt = opts.outfile[-3:]
    if fmt not in ("png", "jpg", "dot", "svg"):
        print >>sys.stderr, \
                "Unexpected file suffix (%s) and no format given, using %s" % \
                (fmt, "dot")
        fmt = "dot"
    file = opts.outfile
elif opts.format and not opts.outfile:
    fmt = opts.format
    file = None
else:
    fmt="dot"
    file = None

# Get the topology from a file or an experiment
if not opts.file:
    if not cert:
        sys.exit("Couldn't find an identity certificate, RPC (Info) needed.")
    try:
        top = get_experiment_topo(opts, cert, url)
    except RPCException, e:
        exit_with_fault(e)
    except RuntimeError, e:
        sys.exit("%s" % e)
else:
    if opts.exp_name or opts.exp_certfile:
        print >>sys.stderr, \
                "Ignoring experiment name or certificate, reading %s" \
                % opts.file
    try:
        top = get_file_topo(opts, cert, url)
    except RPCException, e:
        print >>sys.stderr, "Cannot extract a topology from %s" % opts.file
        exit_with_fault(e)
    except RuntimeError, e:
        sys.exit("Cannot extract a topology from %s: %s" % (opts.file, e))
    except EnvironmentError, e:
        sys.exit("Cannot extract a topology from %s: %s" % (opts.file, e))

    if not top:
        sys.exit("Cannot extract a topology from %s" % opts.file)

# and make an image
if not opts.serialize_only: 
    try:
        if gen_image(top, file, fmt, opts.neato, opts.labels, opts.pixels,
                opts.group) !=0:
            sys.exit("Error rendering graph (subcommand returned non-0)")
    except EnvironmentError, e:
        sys.exit("Failed to draw graph: %s" %e)
    except RuntimeError, e:
        sys.exit("Failed to draw graph: %s" %e)
else:
    sys.exit(0)