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

#!/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
from multiprocessing import Process, Manager, Pool, Queue, freeze_support

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger()



DATASTORE='store.yaml'

class ConnectError(Exception):
  pass

def get_snmp_stats(logger, 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(logger, 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()

def find_right_snmp_ip(logger, data):
  for k,v in data.iteritems():
    if k.startswith('iface_'):
      ip = v['ip'].split('/')[0]
      logger.info("Trying ip %s", ip)
      try:
        uptime = get_snmp_value(logger, ip, 'DISMAN-EVENT-MIB::sysUpTimeInstance')
        return ip
      except ConnectError:
        pass
  return None

    

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

def process_file(logger, store, nf, rescan):
  data = yaml.load(open(nf,'r'))
  nodename = data['nodename']

  if store['snmp'].has_key(nodename):
    ip = store['snmp'][nodename]
    if not ip and rescan:
      logger.info("Re-scanning for new valid IP")
      ip = find_right_snmp_ip(logger, data)
      store['snmp'][nodename] = ip
    else:
      try:
        uptime = get_snmp_value(logger, ip, 'DISMAN-EVENT-MIB::sysUpTimeInstance')
      except ConnectError:
        logger.info("Re-scanning for new valid IP")
        ip = find_right_snmp_ip(logger, data)
        store['snmp'][nodename] = ip
      
  else:
    logger.info("Running discovery for %s", nodename)
    ip = find_right_snmp_ip(logger, data)
    store['snmp'][nodename] = ip

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

    uptime = get_snmp_value(logger, 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

def worker(i, input, m_store):
  logger = logging.getLogger('Worker%s' % i)
  logger.info("Worker")
  for (nf, rescan) in iter(input.get, 'STOP'):
    process_file(logger, m_store, nf, rescan)
  logger.info("END")

if __name__ == '__main__':
  freeze_support()
  task_queue = Queue()
  manager = Manager()

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

  NUMBER_OF_PROCESSES = 10
  RESCAN = True
  plist = {}
  for i in range(NUMBER_OF_PROCESSES):
    plist[i] = Process(target=worker, args=(i, task_queue,m_store))
    plist[i].start()

  for nf in sorted(glob.glob('nodes/*%s*/wleiden.yaml' % ff)):
    task_queue.put((nf, RESCAN))
 
  for i in range(NUMBER_OF_PROCESSES):
    task_queue.put('STOP')

  for i in range(NUMBER_OF_PROCESSES):
    plist[i].join()


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