calculate traffic by snmpwalk for mrtg

#!/usr/local/bin/python
# coding: utf-8
import os
import sys
import time

device = sys.argv[1]
index = sys.argv[2]
community = ''
version = '1'
limit = 86400
interval = 60
output_dir = os.path.join(os.getcwd(), 'output')
output_file = os.path.join(output_dir, '%s_%s.txt' % (device, index))


def is_args():
    if not device and not index:
        print 'Error, wrong parameter'


def seconds():
    return int(time.time())


def traffic_strip(traffic):
    try:
        value = traffic.strip('\n').split(':')[3].strip()
        return value
    except Exception as e:
        print 'Do not have traffic value'
        return ''


def calc_traffic():
    in_cmd = 'snmpwalk -c %s -v %s %s .interfaces.ifTable.ifEntry.ifInOctets.%s' % (community, version, device, index)
    out_cmd = 'snmpwalk -c %s -v %s %s .interfaces.ifTable.ifEntry.ifOutOctets.%s' % (community, version, device, index)
    in_1 = os.popen(in_cmd).read()
    out = os.popen(out_cmd).read()

    result_in = traffic_strip(in_1)
    result_out = traffic_strip(out)
    return result_in, result_out


def calc_real(new, old):
    real = (int(new) - int(old)) / interval
    return real


def write_file(msg):
    with open(output_file, 'a') as f:
        f.write(msg + '\n')


def write_title():
    if not os.path.isdir(output_dir):
        os.mkdir(output_dir)

    with open(output_file, 'w') as title_f:
        title = '%-15s%-15s%-15s%-15s%-15s' % ('TimeStamp', 'OctIn', 'OctOut', 'RealIn', 'RealOut')
        title_f.write(title + '\n')


def main():
    is_args()
    write_title()
    result_in_old, result_out_old = calc_traffic()

    title = '%-15s%-15s%-15s%-15s%-15s' % ('TimeStamp', 'OctIn', 'OctOut', 'RealIn', 'RealOut')
    print title
    limit_bak = limit
    while limit_bak > 0:
        start_time = seconds()
        limit_bak -= 1
        current_time = seconds()

        result_in, result_out = calc_traffic()
        real_in = calc_real(result_in, result_in_old)
        real_out = calc_real(result_out, result_out_old)
        echo = '%-15s%-15s%-15s%-15s%-15s' % (current_time, result_in, result_out, real_in, real_out)
        write_file(echo)
        print echo
        result_in_old = result_in
        result_out_old = result_out

        end_time = seconds()
        dur_time = end_time - start_time
        if dur_time < interval:
            time.sleep(interval - dur_time)


if __name__ == '__main__':
    main()