#!/usr/bin/env python
# vim:ts=2:et:sw=2:ai
#
# Build topological network graph
# Rick van der Zwet <info@rickvanderzwet.nl>
import gformat
import glob
import json
import os
import re
import subprocess
import sys
import tempfile
import yaml


try:
  store = yaml.load(open('store.yaml','r'))
except IOError:
  store = None


HEADER = '''<?xml version="1.0" encoding="UTF-8"?>
<kml xmlns="http://www.opengis.net/kml/2.2">
<Document>
  <Style id="linkDown">
    <LineStyle>
      <color>990000ff</color>
      <width>2</width>
    </LineStyle>
  </Style>
  <Style id="linkUp">
    <LineStyle>
      <color>9900ff00</color>
      <width>3</width>
    </LineStyle>
  </Style>
  <Style id="linkPlanned">
    <LineStyle>
      <color>99ff0000</color>
      <width>1</width>
    </LineStyle>
  </Style>

'''

POINT = '''
  <Placemark>
    <name>%(nodename)s</name>
    <description>%(nodename)s</description>
    <Point>
      <coordinates>%(longitude)s,%(latitude)s,0</coordinates>
    </Point>
  </Placemark>
'''

LINE = '''
   <Placemark>
      <styleUrl>#%(styleId)s</styleUrl>
      <LineString>
        <coordinates>
          %(longitudeA)s,%(latitudeA)s,0
          %(longitudeB)s,%(latitudeB)s,0
        </coordinates>
      </LineString>
    </Placemark>
'''

FOOTER = '''
</Document>
</kml>
'''

def get_graph_data(debug=False):
  poel = {}
  link_type = {}
  link_data = {}
  link_status = {}
  hosts = {}

  try:
    for host in gformat.get_hostlist():
      if debug: print "## Processing host", host
      datadump = gformat.get_yaml(host)
      nodename = datadump['nodename']
      hosts[nodename] = datadump

      for iface_key in datadump['autogen_iface_keys']:
        if not 'ip' in datadump[iface_key]:
            continue

        l = datadump[iface_key]['ip']
        addr, mask = l.split('/')

        # Not parsing of these folks please
        if not gformat.valid_addr(addr):
          continue

        addr = gformat.parseaddr(addr)
        mask = int(mask)
        addr = addr & ~((1 << (32 - mask)) - 1)
        if poel.has_key(addr):
          poel[addr] += [nodename]
          if link_status[addr] == 'linkUp':
            if datadump['status'] == 'planned':
              link_status[addr] = 'linkPlanned'
            elif datadump[iface_key]['status'] == 'planned':
              link_status[addr] = 'linkPlanned'
            elif datadump[iface_key]['status'] == 'down':
              link_status[addr] = 'linkDown'
        else:
          poel[addr] = [nodename]
          # Assume all ns_ip to be 11a for a moment
          if datadump[iface_key].has_key('ns_ip'):
            link_type[addr] = '11a'
          else:
            link_type[addr] = datadump[iface_key]['type']

          if datadump['status'] == 'planned':
            link_status[addr] = 'linkPlanned'
          elif datadump[iface_key]['status'] == 'planned':
            link_status[addr] = 'linkPlanned'
          elif datadump[iface_key]['status'] == 'down':
            link_status[addr] = 'linkDown'
          else:
            link_status[addr] = 'linkUp'

          link_data[addr] = 1
          iface = datadump[iface_key]['autogen_ifname']
          nodename = datadump['nodename']
          INTERVAL = 60 * 10
          if store and store['uptime'].has_key(nodename) and store['snmp'].has_key(nodename) and store['traffic'].has_key(nodename):
            if store and store['traffic'][nodename].has_key(iface):
              (b_in, b_out) = store['traffic'][nodename][iface]
              uptime = store['uptime'][nodename]
              t_kb = float(b_in + b_out) / 1024
              if debug: print "# INFO: Found %s kB in %s seconds" % (t_kb, INTERVAL)
              retval = ((t_kb) / uptime) * INTERVAL
              link_data[addr] = retval

          if debug: print "### %s [%s] is of type %s" % (gformat.showaddr(addr), iface_key, link_type[addr])
  except (KeyError, ValueError), e:
    if debug: print "[FOUT] in '%s' interface '%s'" % (host,iface_key)
    if debug: print e
    sys.exit(1)
  return (poel, link_type, link_data, link_status, hosts)


def make_nodeplanner_json(debug=False):
  # Input data
  poel, link_type, link_data, link_status, hosts = get_graph_data(debug)

  # Export data
  points = []
  for nodename, datadump in hosts.iteritems():
    points.append({'name':  datadump['nodename'], 'lat' : float(datadump['latitude']), 'lon' : float(datadump['longitude']), 'state' : datadump['status']})
  links = []
  for addr,leden in poel.iteritems():
    for index,lid in enumerate(leden[:-1]):
      for buur in leden[index + 1:]:
        links.append({'master' : hosts[buur]['nodename'], 'slave': hosts[lid]['nodename'], 'state' : link_status[addr][4:].lower()})

  # Format data
  return json.dumps({'points' : points, 'links' : links})
    
  

def make_graph(debug=False):
  poel, link_type, link_data, link_status, hosts = get_graph_data(debug)
  output = ''
  if debug: print "# Building KML file"
  output += HEADER
  for nodename, datadump in hosts.iteritems():
    output += POINT % datadump

  for addr,leden in poel.iteritems():
    if link_type[addr] == '11a':
      color = 'green'
      weight = 4
    elif link_type[addr] == 'eth':
      color = 'blue'
      weight = 8
    else:
      color = 'black'
      weight = 1
    width = max(1,link_data[addr])
    leden = sorted(set(leden))
    for index,lid in enumerate(leden[:-1]):
      for buur in leden[index + 1:]:
        keys = {}
        keys['styleId'] = link_status[addr]
        for key, value in hosts[buur].iteritems():
           keys[key + 'A'] = value
        for key, value in hosts[lid].iteritems():
           keys[key + 'B'] = value
        output += LINE % keys
  output += FOOTER
  return output

if __name__ == "__main__":
  OUTFILE = os.path.join(os.getcwd(),'network.kml')
  kml_data = make_graph(debug=True)
  f = open(OUTFILE,'w')
  f.write(kml_data)
  f.close()
  print "# COMPLETED find your output in %s\n" % OUTFILE