流量监控dashboard

#!/usr/bin/env python
# -*- encoding: utf-8 -*-

import os, yaml, time
import cPickle as pickle
import threading
from pysnmp.entity.rfc3413.oneliner import cmdgen


def get_config_info(file):
    # 获取配置文件中的switch信息
    with open(file, 'r') as f:
        content = yaml.load(f)

    return content['web'], content['interval'], content['switch']

def mib_vars(mib, oids, indices = None):
    if indices is None:
        return [cmdgen.MibVariable(mib, x)  for x in oids.split()]
    else:
        return [cmdgen.MibVariable(mib, x, indices)  for x in oids.split()]

def int_str(x):
    try:
        return int(x)
    except ValueError:
        return str(x)

def get_traffic_snmp(ip, community, interface, *args):
    errorIndication, errorStatus, errorIndex, varBindTable = cmdgen.CommandGenerator().nextCmd(
        cmdgen.CommunityData(community),
        cmdgen.UdpTransportTarget((ip, 161)),
        *args, lookupNames = True, lookupValues = True
    )

    for varBindRow in varBindTable:
        row = [ int_str(val) for name, val in varBindRow if name]
        if row[0] == interface:
            return row[1], row[2]

class SwitchTraffic(threading.Thread):
    def __init__(self, ip, community, interface):
        threading.Thread.__init__(self)
        self.ip = ip
        self.community = community
        self.interface = interface

    def run(self):
        oids = mib_vars('IF-MIB', 'ifName ifHCInOctets ifHCOutOctets')
        file = os.path.join('/tmp', 'cache-' + self.ip)
    
        if os.path.exists(file):
            with open(file, 'rb+') as f:
                # 读上一次流量信息
                p = pickle.load(f)
                time_pre, in_pre, out_pre = (p[0], p[1], p[2])
            
                # 获取当前流量信息，写入cache文件
                in_cur, out_cur = get_traffic_snmp(self.ip, self.community, self.interface, *oids)
                time_cur = int(time.time())
                pickle.dump([time_cur, in_cur, out_cur], f)

            # 计算2次流量的变化值，做为当前流量
            if in_cur - in_pre != 0:
                total = (in_cur * 8 - in_pre * 8)
                diff = time_cur - time_pre if time_cur - time_pre != 0 else 0
                in_mbit = float(total) / diff / 1000 / 1000
            else:
                in_mbit = 0

            if out_cur - out_pre != 0:
                total = (out_cur * 8 - out_pre * 8)
                diff = time_cur - time_pre if time_cur - time_pre != 0 else 0
                out_mbit = float(total) / diff / 1000 / 1000
            else:
                out_mbit = 0

            #return round(in_mbit, 2), round(out_mbit, 2)
            print round(in_mbit, 2), round(out_mbit, 2)

        else:
            # 第1次，建立cache文件，记录流量信息，返回当前流量值0
            with open(file, 'wb') as f:
                in_pre, out_pre = get_traffic_snmp(self.ip, self.community, self.interface, *oids)
                time_pre = int(time.time())
                pickle.dump([time_pre, in_pre, out_pre], f)

            #return 0, 0
            print 0, 0

def main():
    file = 'dashboard.yaml'
    web, interval, switch = get_config_info(file)

    for i in switch:
        thread = SwitchTraffic(i['ip'], i['community'], i['interface'])
        thread.start()

if __name__ == '__main__':
    main()