# $Id$
# $URL$

"""
Scan the VINI Central database and create topology "rspec" tags for
slices that have an EGRE key.  This script to be run from a cron job.
"""

import string
import socket
from topology import links, bwlimits

class Node:
    def __init__(self, node):
        self.id = node['node_id']
        self.hostname = node['hostname']
        self.shortname = self.hostname.replace('.vini-veritas.net', '')
        self.site_id = node['site_id']
        self.ipaddr = socket.gethostbyname(self.hostname)

    def get_link_id(self, remote):
        if self.id < remote.id:
            link = (self.id<<7) + remote.id
        else:
            link = (remote.id<<7) + self.id
        return link
        
    def get_iface_id(self, remote):
        if self.id < remote.id:
            iface = 1
        else:
            iface = 2
        return iface
    
    def get_virt_ip(self, remote):
        link = self.get_link_id(remote)
        iface = self.get_iface_id(remote)
        first = link >> 6
        second = ((link & 0x3f)<<2) + iface
        return "192.168.%d.%d" % (first, second)

    def get_virt_net(self, remote):
        link = self.get_link_id(remote)
        first = link >> 6
        second = (link & 0x3f)<<2
        return "192.168.%d.%d/30" % (first, second)
        
    def get_site(self, sites):
        return sites[self.site_id]
            
    
    # What nodes in the set of node_ids are adjacent to this one?
    def adjacent_nodes(self, sites, nodes, node_ids):
        mysite = self.get_site(sites)
        adj_ids = mysite.adj_node_ids.intersection(node_ids)
        adj_nodes = []
        for id in adj_ids:
            adj_nodes.append(nodes[id])
        return adj_nodes
    
    def init_rspecs(self):
        self.rspecs = []
        
    def add_rspec(self, remote, bw):
        my_ip = self.get_virt_ip(remote)
        remote_ip = remote.get_virt_ip(self)
        net = self.get_virt_net(remote)
        rspec = remote.id, remote.ipaddr, bw, my_ip, remote_ip, net
        self.rspecs.append(rspec)

        
class Site:
    def __init__(self, site):
        self.id = site['site_id']
        self.node_ids = site['node_ids']
        self.adj_site_ids = set()
        self.adj_node_ids = set()

    def get_sitenodes(self, nodes):
        n = []
        for i in self.node_ids:
            n.append(nodes[i])
        return n
    
    def add_adjacency(self, site):
        self.adj_site_ids.add(site.id)
        for n in site.node_ids:
            self.adj_node_ids.add(n)
        
    
class Slice:
    def __init__(self, slice):
        self.id = slice['slice_id']
        self.name = slice['name']
        self.node_ids = set(slice['node_ids'])
        self.slice_tag_ids = slice['slice_tag_ids']
    
    def get_tag(self, tagname, slicetags, node = None):
        for i in self.slice_tag_ids:
            tag = slicetags[i]
            if tag.tagname == tagname:
                if (not node) or (node.id == tag.node_id):
                    return tag
        else:
            return None
        
    def get_nodes(self, nodes):
        n = []
        for id in self.node_ids:
            n.append(nodes[id])
        return n
             
    
    # Add a new slice tag   
    def add_tag(self, tagname, value, slicetags, node = None):
        record = {'slice_tag_id':None, 'slice_id':self.id, 'tagname':tagname, 'value':value}
        if node:
            record['node_id'] = node.id
        else:
            record['node_id'] = None
        tag = Slicetag(record)
        slicetags[tag.id] = tag
        self.slice_tag_ids.append(tag.id)
        tag.changed = True       
        tag.updated = True
        return tag
    
    # Update a slice tag if it exists, else add it             
    def update_tag(self, tagname, value, slicetags, node = None):
        tag = self.get_tag(tagname, slicetags, node)
        if tag and tag.value == value:
            value = "no change"
        elif tag:
            tag.value = value
            tag.changed = True
        else:
            tag = self.add_tag(tagname, value, slicetags, node)
        tag.updated = True
            
    def assign_egre_key(self, slicetags):
        if not self.get_tag('egre_key', slicetags):
            try:
                key = free_egre_key(slicetags)
                self.update_tag('egre_key', key, slicetags)
            except:
                # Should handle this case...
                pass
        return
            
    def turn_on_netns(self, slicetags):
        tag = self.get_tag('netns', slicetags)
        if (not tag) or (tag.value != '1'):
            self.update_tag('netns', '1', slicetags)
        return
   
    def turn_off_netns(self, slicetags):
        tag = self.get_tag('netns', slicetags)
        if tag and (tag.value != '0'):
            tag.delete()
        return
    
    def add_cap_net_admin(self, slicetags):
        tag = self.get_tag('capabilities', slicetags)
        if tag:
            caps = tag.value.split(',')
            for cap in caps:
                if cap == "CAP_NET_ADMIN":
                    return
            else:
                newcaps = "CAP_NET_ADMIN," + tag.value
                self.update_tag('capabilities', newcaps, slicetags)
        else:
            self.add_tag('capabilities', 'CAP_NET_ADMIN', slicetags)
        return
    
    def remove_cap_net_admin(self, slicetags):
        tag = self.get_tag('capabilities', slicetags)
        if tag:
            caps = tag.value.split(',')
            newcaps = []
            for cap in caps:
                if cap != "CAP_NET_ADMIN":
                    newcaps.append(cap)
            if newcaps:
                value = ','.join(newcaps)
                self.update_tag('capabilities', value, slicetags)
            else:
                tag.delete()
        return

    # Update the vsys/setup-link and vsys/setup-nat slice tags.
    def add_vsys_tags(self, slicetags):
        link = nat = False
        for i in self.slice_tag_ids:
            tag = slicetags[i]
            if tag.tagname == 'vsys':
                if tag.value == 'setup-link':
                    link = True
                elif tag.value == 'setup-nat':
                    nat = True
        if not link:
            self.add_tag('vsys', 'setup-link', slicetags)
        if not nat:
            self.add_tag('vsys', 'setup-nat', slicetags)
        return


class Slicetag:
    newid = -1 
    def __init__(self, tag):
        self.id = tag['slice_tag_id']
        if not self.id:
            # Make one up for the time being...
            self.id = Slicetag.newid
            Slicetag.newid -= 1
        self.slice_id = tag['slice_id']
        self.tagname = tag['tagname']
        self.value = tag['value']
        self.node_id = tag['node_id']
        self.updated = False
        self.changed = False
        self.deleted = False
    
    # Mark a tag as deleted
    def delete(self):
        self.deleted = True
        self.updated = True
    
    def write(self, slices, nodes, dryrun):
        if not dryrun:
            if self.changed:
                if int(self.id) > 0:
                    UpdateSliceTag(self.id, self.value)
                else:
                    AddSliceTag(self.slice_id, self.tagname, self.value, self.node_id)
            elif self.deleted and int(self.id) > 0:
                try:
                    DeleteSliceTag(self.id)
                except:
                    print "[%s] %s: could not delete" % (self.id, self.tagname)
        else:
            try:
                slice = slices[self.slice_id].name
            except:
                return
            if self.node_id:
                node = nodes[tag.node_id].hostname
            if self.updated:
                if self.deleted:
                    self.value = "deleted"
                elif not self.changed:
                    self.value = "no change"
                if int(self.id) < 0:
                    self.id = "new"
                if self.node_id:
                    print "[%s] %s: %s (%s, %s)" % (self.id, self.tagname, self.value, slice, node)
                else:
                    print "[%s] %s: %s (%s)" % (self.id, self.tagname, self.value, slice)


"""
Create a dictionary of site objects keyed by site ID
"""
def get_sites():
    tmp = []
    for site in GetSites():
        t = site['site_id'], Site(site)
        tmp.append(t)
    return dict(tmp)


"""
Create a dictionary of node objects keyed by node ID
"""
def get_nodes():
    tmp = []
    for node in GetNodes():
        t = node['node_id'], Node(node)
        tmp.append(t)
    return dict(tmp)

"""
Create a dictionary of slice objects keyed by slice ID
"""
def get_slices():
    tmp = []
    for slice in GetSlices():
        t = slice['slice_id'], Slice(slice)
        tmp.append(t)
    return dict(tmp)

"""
Create a dictionary of slicetag objects keyed by slice tag ID
"""
def get_slice_tags():
    tmp = []
    for tag in GetSliceTags():
        t = tag['slice_tag_id'], Slicetag(tag)
        tmp.append(t)
    return dict(tmp)
    
"""
Find a free EGRE key
"""
def free_egre_key(slicetags):
    used = set()
    for i in slicetags:
        tag = slicetags[i]
        if tag.tagname == 'egre_key':
            used.add(int(tag.value))
                
    for i in range(1, 256):
        if i not in used:
            key = i
            break
    else:
        raise KeyError("No more EGRE keys available")
        
    return "%s" % key
   
# For debugging
dryrun = 0

sites = get_sites()
nodes = get_nodes()
slices = get_slices()
slicetags = get_slice_tags()

# Add adjacencies
for (a, b) in links:
    sites[a].add_adjacency(sites[b])
    sites[b].add_adjacency(sites[a])  

for i in slices:
    slice = slices[i]
    tag = slice.get_tag('vini_topo', slicetags)
    if tag:
        topo_type = tag.value
    else:
        topo_type = None
    
    """
    Valid values of topo_type: 
    'vsys':   Use vsys topology scripts to set up virtual links
    'iias':   Automatically create a virtual topology mirroring the physical one
    'manual': Don't modify the topo_rspec tags if present
    None:     No virtual topology
    """
    if topo_type in ['vsys', 'iias', 'manual']:
        slice.assign_egre_key(slicetags)
        slice.turn_on_netns(slicetags)
        slice.add_cap_net_admin(slicetags)
    else:
        # Let them keep EGRE key for now...
        slice.turn_off_netns(slicetags)
        slice.remove_cap_net_admin(slicetags)
            
    # Add vsys/setup-link and vsys/setup-nat
    if topo_type == 'vsys' and slice.get_tag('egre_key', slicetags):
        slice.add_vsys_tags(slicetags)
        
    if topo_type == 'iias' and slice.get_tag('egre_key', slicetags):
        if dryrun:
            print "Building virtual topology for %s" % slice.name
            
        if slice.name in bwlimits:
            bw = bwlimits[slice.name]
        else:
            bw = "1Mbit"

        hosts = "127.0.0.1\t\tlocalhost\n"
        """
        For each node in the slice, check whether the slice is running on any
        adjacent nodes.  For each pair of adjacent nodes, add to nodes' rspecs.
        """
        for node in slice.get_nodes(nodes):
            node.init_rspecs()
            adj_nodes = node.adjacent_nodes(sites, nodes, slice.node_ids)
            for adj in adj_nodes:
                node.add_rspec(adj, bw)
                hosts += "%s\t\t%s\n" % (node.get_virt_ip(adj), node.shortname)
            if node.rspecs:
                topo_str = "%s" % node.rspecs
                slice.update_tag('topo_rspec', topo_str, slicetags, node)
                    
        slice.update_tag('hosts', hosts, slicetags)
    else:
        if dryrun:
            print "Slice %s not using IIAS" % slice.name

    if topo_type == 'manual' and slice.get_tag('egre_key', slicetags):
        for node in slice.get_nodes(nodes):
            topo_tag = slice.get_tag('topo_rspec', slicetags, node)
            if topo_tag:
                topo_tag.updated = True
            
# Update the tag values in the database
for i in slicetags:
    tag = slicetags[i]
    if (tag.tagname == 'topo_rspec' or tag.tagname == 'hosts') and not tag.updated:
        tag.delete()
    tag.write(slices, nodes, dryrun)