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

#!/usr/bin/env python
#
# XXX: Parsing snmpwalk is soo wrong todo, use a proper python library.
#
#
# Rick van der Zwet <info@rickvanderzwet.nl>
#
import gformat
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, m_snmp, m_traffic, m_uptime, host, rescan):
  data = gformat.get_yaml(host)
  nodename = data['nodename']

  if m_snmp.has_key(nodename):
    ip = m_snmp[nodename]
    if not ip and rescan:
      logger.info("Re-scanning for new valid IP")
      ip = find_right_snmp_ip(logger, data)
      m_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)
        m_snmp[nodename] = ip
  else:
    logger.info("Running discovery for %s", nodename)
    ip = find_right_snmp_ip(logger, data)
    m_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')
    m_uptime[nodename] = int(uptime)

    traffic = {}
    for i,f in iface.iteritems():
      traffic[f] = (int(ifin[i]), int(ifout[i]))
    m_traffic[nodename] = traffic
  except ConnectError:
    logger.error("Unable to get all data")
    pass

def worker(i, input, m_snmp, m_traffic, m_uptime):
  logger = logging.getLogger('Worker%s' % i)
  logger.info("Worker")
  for (host, rescan) in iter(input.get, 'STOP'):
    process_file(logger, m_snmp, m_traffic, m_uptime, host, 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
  # XXX: Manager.dict has bug of handling dicts inside dicts, using awefull quick
  # XXX: http://bugs.python.org/issue6766
  m_snmp = manager.dict(store['snmp'])
  m_traffic = manager.dict(store['traffic'])
  m_uptime = manager.dict(store['uptime'])

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

  for host in sorted(gformat.get_hostlist()):
    if ff and not ff in host:
      continue
    task_queue.put((host, RESCAN))

  for i in range(NUMBER_OF_PROCESSES):
    task_queue.put('STOP')

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

  store = {'snmp' : dict(m_snmp), 'traffic' : dict(m_traffic), 'uptime' : dict(m_uptime)}
  yaml.dump(store,open(DATASTORE,'w'))
  logger.info("Stored data at %s" % DATASTORE)