source: genesis/tools/get-link-traffic.py@ 9987

#!/usr/bin/env python
#
# XXX: Parsing snmpwalk is soo wrong todo, use a proper python library.
#
# Rick van der Zwet <info@rickvanderzwet.nl>
#
import glob
import logging
import subprocess
import sys
import yaml
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger()

DATASTORE='store.yaml'

try:
  store = yaml.load(open(DATASTORE,'r'))
except IOError:
  store = { 'snmp' : {}, 'traffic' : {}, 'uptime' : {}}
  pass

class ConnectError(Exception):
  pass

def find_right_snmp_ip(data):
  for k,v in data.iteritems():
    if k.startswith('iface_'):
      ip = v['ip'].split('/')[0]
      logger.info("Trying ip %s", ip)
      p = subprocess.Popen("snmpwalk -t 1 -r 1 -Oq -c public -v2c %s uptime" % ip,
        shell=True,
        stdout=subprocess.PIPE, stderr=subprocess.PIPE)
      (stdout, stderr) = p.communicate()
      if p.returncode == 0:
        return ip
      else:
        logger.error(stderr.strip())
  raise None

def get_snmp_stats(ip,target):
  p = subprocess.Popen("snmpwalk -t 1 -r 1 -Oq -c public -v2c %s %s" % (ip, target), 
    shell=True,
    stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  (stdout, stderr) = p.communicate()
  if p.returncode != 0:
    logger.error(stderr.strip())
    raise ConnectError
  r = {}
  for line in stdout.strip().split('\n'):
    index = line.split()[0][len(target)+1:]
    value = line.split()[1]
    r[index] = value
  return r

def get_snmp_value(ip, target):
  p = subprocess.Popen("snmpget -t 1 -r 1 -Ovt -c public -v2c %s %s" % (ip, target), 
    shell=True,
    stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  (stdout, stderr) = p.communicate()
  if p.returncode != 0:
    logger.error(stderr.strip())
    raise ConnectError
  return stdout.strip()

    

try:
  ff = sys.argv[1]
except IndexError:
  ff = ''

try:
  for nf in sorted(glob.glob('nodes/*%s*/wleiden.yaml' % ff)):
    data = yaml.load(open(nf,'r'))
    nodename = data['nodename']
    
    if store['snmp'].has_key(nodename):
      ip = store['snmp'][nodename]
    else:
      logger.info("Running discovery for %s", nodename)
      ip = find_right_snmp_ip(data)
      store['snmp'][nodename] = ip

    if ip == None:
      logger.error("No valid ip found for node %s", nodename)
      continue
    
    logger.info("Processing %s via %s", nodename, ip)
    target = 'IF-MIB::ifDescr'
    
    try:
      iface = get_snmp_stats(ip, 'IF-MIB::ifDescr')
      ifout = get_snmp_stats(ip, 'IF-MIB::ifOutOctets')
      ifin = get_snmp_stats(ip, 'IF-MIB::ifInOctets')

      uptime = get_snmp_value(ip, 'DISMAN-EVENT-MIB::sysUpTimeInstance')
      store['uptime'][nodename] = int(uptime)
      
      traffic = {}
      for i,f in iface.iteritems():
        traffic[f] = (int(ifin[i]), int(ifout[i]))
      store['traffic'][nodename] = traffic 
    except ConnectError:
      logger.error("Unable to get all data")
      pass
except KeyboardInterrupt:
  pass

yaml.dump(store,open(DATASTORE,'w'))