#!/usr/bin/env python
## Copyright (c) 2004, The Regents of the University of California, through
## Lawrence Berkeley National Laboratory (subject to receipt of any required
## approvals from the U.S. Dept. of Energy).  All rights reserved.
"""
Parse vmstat output to get CPU usage
"""
__author__ = "Dan Gunter"
__created__ = "5 November 2004"
__rcsid__ = "$Id: nl_cpu 31131 2012-04-10 12:38:18Z dang $"

import logging
import signal
import os
import sys
import time
from subprocess import *


from netlogger.nllog import OptionParser, get_logger
from netlogger import nlapi, nllog

## Globals

PREFIX = "cpu"
EVENTS = (PREFIX+".usr", PREFIX+".sys", PREFIX+".idle")
OS = os.uname()[0]

## Functions

def vmstatCpu(log, fields, max_iter, cmd):
    """vmstatCpu(log:nllite.LogOutputStream,fields:list[bool]) -> None

    Parse vmstat for CPU utilization.
    'fields' tells whether user, sys, idle should be logged
    """
    iter = 0
    #print "calling: ", cmd
    p = Popen(cmd, shell=True, stdout=PIPE )
    f = p.stdout
    #f = os.popen("/usr/bin/vmstat 1")
    # skip over header lines
    f.readline()
    f.readline()
    f.readline() # skip first set of actual output too, often not accurate, or even negative
    while 1:
        line = f.readline()
        #print line, OS
        try:
            if OS == 'Linux':
                vals = map(int, line.split()[-4:-1])
            else:
                vals = map(float, line.split()[-3:])
        except (ValueError, IndexError), err:
            #print "ERROR: %s" % err
            continue
        #print fields
        #print vals
        for i in range(3):
            if fields[i]:
                log.info(EVENTS[i], VAL=vals[i])
        iter = iter + 1
	if max_iter != -1 and iter >= max_iter:
	     # kill off vmstat and return
             os.kill(p.pid, signal.SIGKILL)
	     return;

def parse_rot_handler(value):
    if len(value) < 1:
        raise ValueError("Bad log rotation interval, too short")
    tm_unit = value[-1].lower()
    if tm_unit not in ('h', 'm', 'd'):
        raise ValueError("Bad log rotation unit '%s' "
                         "not in m,h,d" % tm_unit)
    try:
        tm_interval = int(value[:-1])
    except ValueError:
        raise ValuerError("Log rotation value '%s' must be an integer" %
                          value[:-1])
    return dict(when=tm_unit, interval=tm_interval)

def main():
    desc = ' '.join(__doc__.split())
    parser = OptionParser(description=desc)
    parser.add_option("-a", "--all", action="store_true", dest="all",
                      help="Report all types of utilization (default)")
    parser.add_option("-b", "--backup", action="store", type="int",
                      dest="backup_count", default=10, metavar="COUNT",
                      help="For -r/--rotate, if this is > 0, then "
                      "when rotation is done, no more than COUNT files "
                      "are kept - the oldest ones are deleted "
                      "(default=%default)")
    parser.add_option("-c", "--count", action="store", dest="count",
                      default=-1, metavar="NUM", type='int',
                      help="Repeat count times (default=forever)")
    parser.add_option("-f", "--logfile", action="store", dest="logfile",
                      default="-", metavar="FILE",
                      help="Output file (default='-', stdout)."
                      " See -r/--rotate for file rotation options")
    parser.add_option("-i", "--idle", action="store_true", dest="idle",
                      help="Report 'idle' utilization")
    parser.add_option("-p", "--program", action="store", dest='vmstat',
                      default="/usr/bin/vmstat",
                      help="vmstat executable (default=%s)")
    parser.add_option('-r','--rotate', default=None, action='store',
                      dest='rotate', metavar='TIME',
                      help="rotate logs at an interval "
                      "(<N>d or <N>h or <N>m). See also -b/--backup")
    parser.add_option("-s", "--sys", action="store_true", dest="sys",
                      help="Report 'system' utilization")
    parser.add_option("-S", "--syslog", action="store_true", dest="syslog",
                      help="Send logs to local syslog, ignores -f/--logfile")
    parser.add_option("-u", "--user", action="store_true", dest="usr",
                      help="Report 'user' utilization")
    parser.add_option("-w", "--wait", action="store", dest="wait",
                      default=1, metavar="SEC", type='int',
                      help="Pause wait seconds between each display. "
                      "(default=1)")
    options, args = parser.parse_args()
    log = get_logger(__file__)  # Should be first done, just after parsing args
    # check for program
    if not os.path.exists(options.vmstat):
        parser.error("vmstat executable '%s' not found" % options.vmstat)
    # figure out which fields to log: (usr, sys, idle)
    if options.all:
        fields = [True] * 3
    else:
        fields = [False] * 3
        if options.usr:
            fields[0] = True
        if options.sys:
            fields[1] = True
        if options.idle:
            fields[2] = True
        if sum(fields) == 0:
            fields = [True] * 3
    # open log
    log.info("open.log.start")
    nllog.setLoggerClass(nllog.BPLogger)
    result_log = nllog.get_logger(".vmstat")
    if options.syslog:
        addr = "/dev/log"
        handler = logging.handlers.SysLogHandler(addr)
    elif options.logfile == '-':
        handler = logging.StreamHandler(sys.stdout)
    else:
        if options.rotate is None:
            handler = logging.FileHandler(options.logfile)
        else:
            try:
                rotate_kw = parse_rot_handler(options.rotate)
            except ValueError, err:
                parser.error("Bad log rotation value: %s" % err)
            rotate_kw['backupCount'] = options.backup_count
            handler = logging.handlers.TimedRotatingFileHandler(
                options.logfile, **rotate_kw)
    result_log.addHandler(handler)
    result_log.set_meta(host=nlapi.get_host())
    result_log.setLevel(logging.INFO)
    log.info("open.log.end", status=0)
    log.info("loop.start")
    # write CPU values
    vmstatCpu(result_log, fields, options.count, cmd="/usr/bin/vmstat 1")
    log.info("loop.end", status=0, num=iter)
    return

if __name__ == "__main__":
    sys.exit(main())
