source: fedd/federation/topdl.py @ e794984

#!/usr/local/bin/python

import re

class base:
    @staticmethod
    def init_class(c, arg):
        if isinstance(arg, dict):
            try:
                return c(**arg)
            except:
                print "%s" % arg
                raise
        elif isinstance(arg, c):
            return arg
        else:
            return None

    @staticmethod
    def make_list(a):
        if isinstance(a, basestring) or isinstance(a, dict): return [ a ]
        elif getattr(a, '__iter__', None): return a
        else: return [ a ]

    def get_attribute(self, key):
        rv = None
        attrs = getattr(self, 'attribute', None)
        if attrs:
            for a in attrs:
                if a.attribute == key:
                    rv = a.value
                    break
        return rv

    def set_attribute(self, key, value):
        attrs = getattr(self, 'attribute', None)
        if attrs is None:
            return
        for a in attrs:
            if a.attribute == key: 
                a.value = value
                break
        else:
            attrs.append(Attribute(key, value))

class ConsistencyError(RuntimeError): pass

class Attribute(base):
    def __init__(self, attribute, value):
        self.attribute = attribute
        self.value = value

    def clone(self):
        return Attribute(attribute=self.attribute, value=self.value)

    def to_dict(self):
        return { 'attribute': self.attribute, 'value': self.value }

class Capacity(base):
    def __init__(self, rate, kind):
        self.rate = float(rate)
        self.kind = kind

    def clone(self):
        return Capacity(rate=self.rate, kind=self.kind)

    def to_dict(self):
        return { 'rate': float(self.rate), 'kind': self.kind }

class Latency(base):
    def __init__(self, time, kind):
        self.time = float(time)
        self.kind = kind

    def clone(self):
        return Latency(time=self.time, kind=self.kind)

    def to_dict(self):
        return { 'time': float(self.time), 'kind': self.kind }

class Substrate(base):
    def __init__(self, name, capacity=None, latency=None, attribute=[]):
        self.name = name
        self.capacity = self.init_class(Capacity, capacity)
        self.latency = self.init_class(Latency, latency)
        self.attribute = [ self.init_class(Attribute, a) \
                for a in self.make_list(attribute) ]
        self.interfaces = [ ]

    def clone(self):
        if self.capacity: c = self.capacity.clone()
        else: c = None

        if self.latency: l = self.latency.clone()
        else: l = None

        return Substrate(name=self.name,
                capacity=c,
                latency=l,
                attribute = [a.clone() for a in self.attribute])

    def to_dict(self):
        rv = { 'name': self.name }
        if self.capacity:
            rv['capacity'] = self.capacity.to_dict()
        if self.latency:
            rv['latency'] = self.latency.to_dict()
        if self.attribute:
            rv['attribute'] = [ a.to_dict() for a in self.attribute ]
        return rv

class CPU(base):
    def __init__(self, type, attribute=[]):
        self.type = type
        self.attribute = [ self.init_class(Attribute, a) for a in \
                self.make_list(attribute) ]

    def clone(self):
        return CPU(type=self.type,
                attribute = [a.clone() for a in self.attribute])

    def to_dict(self):
        rv = { 'type': self.type}
        if self.attribute:
            rv['attribute'] = [ a.to_dict() for a in self.attribute ]
        return rv

class Storage(base):
    def __init__(self, amount, persistence, attribute=[]):
        self.amount = float(amount)
        self.presistence = persistence
        self.attribute = [ self.init_class(Attribute, a) \
                for a in self.make_list(attribute) ]

    def clone(self):
        return Storage(amount=self.amount, persistence=self.persistence, 
                attribute = [a.clone() for a in self.attribute])

    def to_dict(self):
        rv = { 'amount': float(self.amount), 'persistence': self.persistence }
        if self.attribute:
            rv['attribute'] = [ a.to_dict() for a in self.attribute ]
        return rv

class OperatingSystem(base):
    def __init__(self, name=None, version=None, distribution=None,
            distributionversion=None, attribute=[]):
        self.name = name
        self.version = version
        self.distribution = distribution
        self.distributionversion = distributionversion
        self.attribute = [ self.init_class(Attribute, a) \
                for a in self.make_list(attribute) ]

    def clone(self):
        return OperatingSystem(name=self.name,
                version=self.version,
                distribution=self.distribution,
                distributionversion=self.distributionversion,
                attribute = [ a.clone() for a in self.attribute])

    def to_dict(self):
        rv = { }
        if self.name: rv['name'] = self.name
        if self.version: rv['version'] = self.version
        if self.distribution: rv['version'] = self.distribution
        if self.distributionversion: rv['version'] = self.distributionversion
        if self.attribute:
            rv['attribute'] = [ a.to_dict() for a in self.attribute ]
        return rv

class Software(base):
    def __init__(self, location, install=None, attribute=[]):
        self.location = location
        self.install = install
        self.attribute = [ self.init_class(Attribute, a)\
                for a in self.make_list(attribute) ]

    def clone(self):
        return Software(location=self.location, install=self.install, 
                attribute=[a.clone() for a in self.attribute])

    def to_dict(self):
        rv = { 'location': self.location }
        if self.install: rv['install'] = self.install
        if self.attribute:
            rv['attribute'] = [ a.to_dict() for a in self.attribute ]
        return rv

class Interface(base):
    def __init__(self, substrate, capacity=None, latency=None, attribute=[],
            element=None):
        self.substrate = self.make_list(substrate)
        self.capacity = self.init_class(Capacity, capacity)
        self.latency = self.init_class(Latency, latency)
        self.attribute = [ self.init_class(Attribute, a) \
                for a in self.make_list(attribute) ]
        self.element = element 
        self.subs = [ ]

    def clone(self):
        if self.capacity: c = self.capacity.clone()
        else: c = None

        if self.latency: l = self.latency.clone()
        else: l = None

        return Interface(substrate=self.substrate,
                capacity=c, latency=l,
                attribute = [ a.clone() for a in self.attribute])

    def to_dict(self):
        rv = { 'substrate': self.substrate }
        if self.capacity:
            rv['capacity'] = self.capacity.to_dict()
        if self.latency:
            rv['latency'] = self.latency.to_dict()
        if self.attribute:
            rv['attribute'] = [ a.to_dict() for a in self.attribute ]
        return rv

class ID(base):
    def __init__(self, fedid=None, uuid=None, uri=None, localname=None,
            kerberosUsername=None):
        self.fedid=fedid
        self.uuid = uuid
        self.uri = uri
        self.localname = localname
        self.kerberosUsername = kerberosUsername

    def clone(self):
        return ID(self.fedid, self.uuid, self.uri, self.localname,
                self.kerberosUsername)

    def to_dict(self):
        rv = { }
        if self.fedid: rv['fedid'] = self.fedid
        if self.uuid: rv['uuid'] = self.uuid
        if self.uri: rv['uri'] = self.uri
        if self.localname: rv['localname'] = self.localname
        if self.kerberosUsername: rv['kerberosUsername'] = self.kerberosUsername
        return rv

class Computer(base):
    def __init__(self, name=[], cpu=[], os=[], software=[], storage=[],
            interface=[], attribute=[]):
        def assign_element(i):
            i.element = self

        self.name = self.make_list(name)
        self.cpu = [ self.init_class(CPU, c)  for c in self.make_list(cpu) ]
        self.os = [ self.init_class(OperatingSystem, c) \
                for c in self.make_list(os) ]
        self.software = [ self.init_class(Software, c) \
                for c in self.make_list(software) ]
        self.storage = [ self.init_class(Storage, c) \
                for c in self.make_list(storage) ]
        self.interface = [ self.init_class(Interface, c) \
                for c in self.make_list(interface) ]
        self.attribute = [ self.init_class(Attribute, a) \
                for a in self.make_list(attribute) ]
        map(assign_element, self.interface)

    def clone(self):
        return Computer(name=self.name,
                cpu=[x.clone() for x in self.cpu],
                os=[x.clone() for x in self.os],
                software=[x.clone() for x in self.software],
                storage=[x.clone() for x in self.storage],
                interface=[x.clone() for x in self.interface],
                attribute=[x.clone() for x in self.attribute])

    def to_dict(self):
        rv = { }
        if self.name:
            rv['name'] = self.name
        if self.cpu:
            rv['cpu'] = [ c.to_dict() for  c in self.cpu ]
        if self.os:
            rv['os'] = [ o.to_dict() for o in self.os ]
        if self.software:
            rv['software'] = [ s.to_dict() for s in self.software ]
        if self.storage:
            rv['storage'] = [ s.to_dict for s in self.storage ]
        if self.interface:
            rv['interface'] = [ i.to_dict() for i in self.interface ]
        if self.attribute:
            rv['attribute'] = [ i.to_dict() for i in self.attribute ]
        return { 'computer': rv }


class Testbed(base):
    def __init__(self, uri, type, interface=[], attribute=[]):
        self.uri = uri
        self.type = type
        self.interface = [ self.init_class(Interface, c) \
                for c in self.make_list(interface) ]
        self.attribute = [ self.init_class(Attribute, c) \
                for c in self.make_list(attribute) ]

    def clone(self):
        return Testbed(self.uri, self.type,
                interface=[i.clone() for i in self.interface],
                attribute=[a.cone() for a in self.attribute])

    def to_dict(self):
        rv = { }
        if self.uri: rv['uri'] = self.uri
        if self.type: rv['type'] = self.type
        if self.interface:
            rv['interface'] = [ i.to_dict() for i in self.interface]
        if self.attribute:
            rv['attribute'] = [ a.to_dict() for a in self.attribute]
        return { 'testbed': rv }

class Segment(base):
    def __init__(self, id, type, uri, interface=[], attribute=[]):
        self.id = self.init_class(ID, id)
        self.type = type
        self.uri = uri
        self.interface = [ self.init_class(Interface, c) \
                for c in self.make_list(interface) ]
        self.attribute = [ self.init_class(Attribute, c) \
                for c in self.make_list(attribute) ]

    def clone(self):
        return Segment(self.id.clone(), self.type, self.uri, 
                interface=[i.clone() for i in self.interface], 
                attribute=[a.clone() for a in self.attribute])

    def to_dict(self):
        rv = { }
        if self.id: rv['id'] = self.id.to_dict()
        if self.type: rv['type'] = self.type
        if self.uri: rv['uri'] = self.uri
        if self.interface:
            rv['interface'] = [ i.to_dict() for i in self.interface ]
        if self.attribute:
            rv['attribute'] = [ a.to_dict() for a in self.attribute ]
        return { 'segment': rv }


class Other(base):
    def __init__(self, interface=[], attribute=[]):
        self.interface = [ self.init_class(Interface, c) \
                for c in self.make_list(interface) ]
        self.attribute = [ self.init_class(Attribute, c) \
                for c in self.make_list(attribute) ]

    def clone(self):
        return Other(interface=[i.clone() for i in self.interface], 
                attribute=[a.clone() for a in attribute])

    def to_dict(self):
        rv = {}
        if self.interface:
            rv['interface'] = [ i.to_dict() for i in self.interface ]
        if self.attribute:
            rv['attribute'] = [ a.to_dict() for a in self.attribute ]
        return {'other': rv }


class Topology(base):
    @staticmethod
    def init_element(e):
        """
        e should be of the form { typename: args } where args is a dict full of
        the right parameters to initialize the element.  e should have only one
        key, but we walk e's keys in an arbitrary order and instantiate the
        first key we know how to.
        """
        classmap = {
                'computer': Computer,
                'testbed': Testbed,
                'segment': Segment,
                'other': Other,
            }

        if isinstance(e, dict):
            for k in e.keys():
                cl = classmap.get(k, None)
                if cl: return cl(**e[k])
        else:
            return e

    def __init__(self, substrates=[], elements=[]):
        self.substrates = [ self.init_class(Substrate, s) \
                for s in self.make_list(substrates) ]
        self.elements = [ self.init_element(e) \
                for e in self.make_list(elements) ]
        self.incorporate_elements()

    def incorporate_elements(self):

        # Could to this init in one gulp, but we want to look for duplicate
        # substrate names
        substrate_map = { }
        for s in self.substrates:
            s.interfaces = [ ]
            if not substrate_map.has_key(s.name):
                substrate_map[s.name] = s
            else:
                raise ConsistencyError("Duplicate substrate name %s" % s.name)

        for e in self.elements:
            for i in e.interface:
                i.element = e
                i.subs = [ ]
                for sn in i.substrate:
                    # NB, interfaces have substrate names in their substrate
                    # attribute.
                    if substrate_map.has_key(sn):
                        sub = substrate_map[sn]
                        i.subs.append(sub)
                        sub.interfaces.append(i)
                    else:
                        raise ConsistencyError("No such substrate for %s" % sn)

    def clone(self):
        return Topology(substrates=[s.clone() for s in self.substrates], 
                elements=[e.clone() for e in self.elements])


    def make_indices(self):
        sub_index = dict([(s.name, s) for s in self.substrates])
        elem_index = dict([(n, e) for e in self.elements for n in e.name])

    def to_dict(self):
        rv = { }
        if self.substrates:
            rv['substrates'] = [ s.to_dict() for s in self.substrates ]
        if self.elements:
            rv['elements'] = [ s.to_dict() for s in self.elements ]
        return rv

def topology_from_xml(string=None, file=None, filename=None, top="topology"):
    import xml.parsers.expat

    class parser:
        def __init__(self):
            self.stack = [ ]
            self.chars = ""
            self.key = ""
            self.have_chars = False
            self.current = { }
            self.in_cdata = False
        
        def start_element(self, name, attrs):
            self.chars = ""
            self.have_chars = False
            self.key = str(name)
            self.stack.append((self.current, self.key))
            self.current = { }

        def end_element(self, name):
            if self.have_chars:
                self.chars = self.chars.strip()
                if len(self.chars) >0:
                    addit = self.chars
                else:
                    addit = self.current
            else:
                addit = self.current

            parent, key = self.stack.pop()
            if parent.has_key(key):
                if isinstance(parent[key], list):
                    parent[key].append(addit)
                else:
                    parent[key] = [parent[key], addit]
            else:
                parent[key] = addit
            self.current = parent
            self.key = key

            self.chars = ""
            self.have_chars = False

        def char_data(self, data):
            self.have_chars = True
            self.chars += data

    p = parser()
    xp = xml.parsers.expat.ParserCreate()

    xp.StartElementHandler = p.start_element
    xp.EndElementHandler = p.end_element
    xp.CharacterDataHandler = p.char_data

    num_set = len([ x for x in (string, filename, file)\
            if x is not None ])

    if num_set != 1:
        raise RuntimeError("Exactly one one of file, filename and string " + \
                "must be set")
    elif filename:
        f = open(filename, "r")
        xp.ParseFile(f)
        f.close()
    elif file:
        xp.ParseFile(file)
    elif string:
        xp.Parse(string, isfinal=True)
    else:
        return None

    return Topology(**p.current[top])

def topology_to_xml(t, top=None):
    """
    Print the topology as XML.  This is quick and dirty, but should work for
    many purposes.  Convert the topology to a dict and print it recursively.
    """
    from xml.sax.saxutils import escape

    def dict_to_xml(e, top=None):
        if top: rv = "<%s>" % top
        else: rv = ""

        for k in e.keys():
            if isinstance(e[k], (basestring, int, float, long)):
                rv += "<%s>%s</%s>" % (k, escape(e[k]), k)
            elif isinstance(e[k], dict):
                rv += "<%s>%s</%s>" % (k, dict_to_xml(e[k]), k)
            elif getattr(e[k], '__iter__', None):
                for ee in e[k]:
                    if isinstance(ee, dict):
                        rv += "<%s>%s</%s>" % (k, dict_to_xml(ee), k)
                    else:
                        rv += "<%s>%s</%s>" % (k, escape(ee), k)
            else:
                raise ConsistencyError("What is this?? %s %s" % (k, e[k]))
        if top: rv += "</%s>" % top
        return rv

    return dict_to_xml(t.to_dict(), top)


def topology_to_vtopo(t):
    nodes = [ ]
    lans = [ ]

    for eidx, e in enumerate(t.elements):
        if e.name: name = e.name[0]
        else: name = "unnamed_node%d" % eidx
        
        ips = [ ]
        for idx, i in enumerate(e.interface):
            ip = i.get_attribute('ip4_address')
            ips.append(ip)
            port = "%s:%d" % (name, idx)
            for idx, s in enumerate(i.subs):
                bw = 100000
                delay = 0.0
                if s.capacity:
                    bw = s.capacity.rate
                if i.capacity:
                    bw = i.capacity.rate

                if s.latency:
                    delay = s.latency.time
                if i.latency:
                    bw = i.latency.time

                lans.append({
                    'member': port,
                    'vname': s.name,
                    'ip': ip,
                    'vnode': name,
                    'delay': delay,
                    'bandwidth': bw,
                    })
        nodes.append({
            'ips': ":".join(ips),
            'vname': name,
            })

    return { 'node': nodes, 'lan': lans }

def to_tcl_name(n):
    t = re.sub('-(\d+)', '(\1)', n)
    return t

def generate_portal_command_filter(cmd):
    def rv(e):
        s =""
        if isinstance(e, Computer):
            gw = e.get_attribute('portal')
            if gw:
                s = "%s $%s\n" % (cmd, to_tcl_name(e.name[0]))
        return s
    return rv

def generate_portal_image_filter(image):
    def rv(e):
        s =""
        if isinstance(e, Computer):
            gw = e.get_attribute('portal')
            if gw:
                s = "tb-set-node-os $%s %s\n" % (to_tcl_name(e.name[0]), image)
        return s
    return rv

def generate_portal_hardware_filter(type):
    def rv(e):
        s =""
        if isinstance(e, Computer):
            gw = e.get_attribute('portal')
            if gw:
                s = "tb-set-hardware $%s %s\n" % (to_tcl_name(e.name[0]), type)
        return s
    return rv


def topology_to_ns2(t, filters=[], routing="Session"):
    out = """
set ns [new Simulator]
source tb_compat.tcl

"""

    for e in t.elements:
        rpms = ""
        tarfiles = ""
        if isinstance(e, Computer):
            name = to_tcl_name(e.name[0])
            out += "set %s [$ns node]\n" % name
            if e.os and len(e.os) == 1:
                osid = e.os[0].get_attribute('osid')
                if osid:
                    out += "tb-set-node-os $%s %s\n" % (name, osid)
            hw = e.get_attribute('type')
            if hw:
                out += "tb-set-hardware $%s %s\n" % (name, hw)
            for s in e.software:
                if s.install:
                    tarfiles += "%s %s " % (s.install, s.location)
                else:
                    rpms += "%s " % s.location
            if rpms:
                out += "tb-set-node-rpms $%s %s\n" % (name, rpms)
            if tarfiles:
                out += "tb-set-node-tarfiles $%s %s\n" % (name, tarfiles)
            startcmd = e.get_attribute('startup')
            if startcmd:
                out+= 'tb-set-node-startcmd $%s "%s"\n' % (name, startcmd)
            for f in filters:
                out += f(e)
            out+= "\n"
    
    for idx, s in enumerate(t.substrates):
        loss = s.get_attribute('loss')
        if s.latency:
            delay = s.latency.time
        else:
            delay = 0
        name = to_tcl_name(s.name or "sub%d" % idx)

        if len(s.interfaces) > 2:
            # Lan
            members = [ to_tcl_name("$%s") % i.element.name[0] \
                    for i in s.interfaces]
            out += 'set %s [$ns make-lan "%s" %fkb %fms ]\n' % \
                    (name, " ".join(members), s.capacity.rate, delay)
            if loss:
                "tb-set-lan-loss $%s %f\n" % (name, float(loss))

            for i in s.interfaces:
                e = i.element
                ip = i.get_attribute("ip4_address")
                if ip:
                    out += "tb-set-ip-lan $%s $%s %s\n" % \
                            (to_tcl(e.name[0]), name, ip)
                if i.capacity and i.capacity.rate != s.capacity.rate:
                    out += "tb-set-node-lan-bandwidth $%s $%s %fkb\n" % \
                            (to_tcl_name(e.name[0]), name, i.capacity.rate)
                if i.latency and i.latency.time != delay:
                    out += "tb-set-node-lan-delay $%s $%s %fms\n" % \
                            (to_tcl_name(e.name[0]), name, i.latency.time)
                iloss = i.get_attribute('loss')
                if loss and iloss != loss :
                    out += "tb-set-node-lan-loss $%s $%s %f\n" % \
                            (to_tcl_name(e.name[0]), name, float(loss))
            out+= "\n"
        elif len(s.interfaces) == 2:
            f = s.interfaces[0]
            t = s.interfaces[1]

            out += "set %s [$ns duplex-link $%s $%s %fkb %fms DropTail]\n" %\
                    (name, to_tcl_name(f.element.name[0]), 
                            to_tcl_name(t.element.name[0]),
                            s.capacity.rate, delay)
            if loss:
                out += "tb-set-link-loss $%s %f\n" % (name, float(loss))

            for i in s.interfaces:
                lloss = i.get_attribute("loss")
                cap_override = i.capacity and \
                        i.capacity.rate != s.capacity.rate
                delay_override = i.latency and \
                        i.latency.time != delay
                loss_override = lloss and lloss != loss
                if cap_override or delay_override or loss_override:
                    if i.capacity: cap = i.capacity.rate
                    else: cap = s.capacity.rate

                    if i.latency: delay = i.latency.time

                    if lloss: loss = lloss
                    else: loss = loss or 0.0

                    out += "tb-set-link-simplex-params $%s $%s %fms %fkb %f\n"\
                            % (name, to_tcl_name(i.element.name[0]),
                                    delay, cap, loss)
                ip = i.get_attribute('ip4_address')
                if ip:
                    out += "tb-set-ip-link $%s $%s %s\n" % \
                            (to_tcl_name(i.element.name[0]), name, ip)
            out+= "\n"
        for f in filters:
            out+= f(s)
    out+="$ns rtproto %s" % routing
    out+="""
$ns run
"""
    return out