source: fedd/fedd_image.py @ 2c51061

#!/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")

def gen_dot_topo(t, labels, dotfile):
    """
    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]

    # Output the nodes.  Deal with the possibility that there's no name, but
    # that's unlikely.
    for i, n in enumerate(t.elements):
        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]'

    # Encode the lans and the links
    for l in lans:
        print >>dotfile, '\t"%s" [shape=ellipse, style=filled,\\' % l.name
        print >>dotfile,'\t\tcolor=black,fillcolor="#80c0f8",regular=1]'
        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):
    """
    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:
        for f in ['/usr/bin/neato', '/usr/local/bin/neato', 
                '/usr/bin/dot', '/usr/local/bin/dot']:
            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];' \
                % (size, size, dpi)
    else:
        print >>dotfile, '\tgraph [size="%i,%i", ratio=fill];' \
                % (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)
    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):
    """
    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,
            opts.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):
    """
    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, 
            opts.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 = 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)
    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)
    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) !=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)