#!/usr/bin/python
# -*- coding: utf-8 -*-
# vim: set fileencoding=utf-8
#
# Munin plugin to show the network I/O per vm
# On redhat based systems
#
# Copyright Igor Borodikhin
# Copyright Peter Meier
#
# License : GPLv3
#
#
# parsed environment variables:
# vmsuffix: part of vm name to be removed
#
#%# capabilities=autoconf
#%# family=contrib

import re, os, sys
from subprocess import Popen, PIPE

def config(vms):
    ''' Print the plugin's config
    @param vm_names : a list of "cleaned" vms' name
    '''
    base_config = """graph_title KVM Network I/O
graph_vlabel Bytes rx(-)/tx(+) per second
graph_category KVM
graph_info This graph shows the network I/O of the virtual machines
graph_args --base 1024"""
    print base_config
    for pid in vms:
        macs = get_vm_macs(pid)
        i = 0
        for mac in macs:
          print "%s_eth%s_in.label %s_eth%s" % (vms[pid],i, vms[pid], i)
          print "%s_eth%s_in.type COUNTER" % (vms[pid], i)
          print "%s_eth%s_in.min 0" % (vms[pid],i)
          print "%s_eth%s_in.draw LINE2" % (vms[pid],i)
          print "%s_eth%s_out.negative %s_eth%s_in" % (vms[pid], i, vms[pid], i)
          print "%s_eth%s_out.label %s_eth%s" % (vms[pid], i, vms[pid], i)
          print "%s_eth%s_out.type COUNTER" % (vms[pid], i)
          print "%s_eth%s_out.min 0" % (vms[pid], i)
          print "%s_eth%s_out.draw LINE2" % (vms[pid], i)
          i += 1

def clean_vm_name(vm_name):
    ''' Replace all special chars
    @param vm_name : a vm's name
    @return cleaned vm's name
    '''
    # suffix part defined in conf
    suffix = os.getenv('vmsuffix')
    if suffix:
        vm_name = re.sub(suffix,'',vm_name)

    return re.sub(r"[^a-zA-Z0-9_]", "_", vm_name)
    
def fetch(vms):
    ''' Fetch values for a list of pids
    @param dictionnary {kvm_pid: cleaned vm name}
    '''
    res = {}
    macs_to_inf = find_macs_to_inf()
    interfaces = {}
    for e in Popen('cat /proc/net/dev | awk \'{ print $1  ":" $9 }\'', shell=True, stdout=PIPE).communicate()[0].split('\n'):
        s = e.split(':')
        if len(s) == 3:
            interfaces[s[0]] = (s[1],s[2])
    for pid in vms:
        macs = get_vm_macs(pid)
        i = 0
        for mac in macs:
            inf = macs_to_inf[mac]
            values = interfaces[inf]
            if len(values) == 2:
                print "%s_eth%s_in.value %s" % (vms[pid], i, values[0])
                print "%s_eth%s_out.value %s" % (vms[pid], i, values[1])
            i += 1

def detect_kvm():
    ''' Check if kvm is installed
    '''
    kvm = Popen("which kvm", shell=True, stdout=PIPE)
    kvm.communicate()
    return not bool(kvm.returncode)

def find_vm_names(pids):
    '''Find and clean vm names from pids
    @return a dictionnary of {pids : cleaned vm name}
    '''
    result = {}
    for pid in pids:
        cmdline = open("/proc/%s/cmdline" % pid, "r")
        result[pid] = clean_vm_name(re.sub(r"^.*-name\x00([a-zA-Z0-9.-]*)\x00\-.*$",r"\1", cmdline.readline()))        
    return result
    
def get_vm_macs(pid):
    '''Find macs for a pid
    @return the mac addresses for a specified pid
    '''
    cmdline = open("/proc/%s/cmdline" % pid, "r")
    line = cmdline.readline()
    # macs are fe:... on the host
    macs = [ re.sub(r"^\d{2}",'fe',p.split('=')[1]) for p in line.split(",") if re.match(r"^mac(addr)?=",p) ]
    return macs

def list_pids():
    ''' Find the pid of kvm processes
    @return a list of pids from running kvm
    '''
    pid = Popen("pidof qemu-kvm kvm", shell=True, stdout=PIPE)
    return pid.communicate()[0].split()

def find_macs_to_inf():
    ''' Find interfaces for vms
    @return a dictionary of macs to inf
    '''
    result = {}
    inf = ""
    kvm = Popen("ip a | grep -E -A 1 '(tap|vnet)' | awk '{print $2}' | grep -v '^$'", shell=True, stdout=PIPE)
    res = kvm.communicate()[0].split('\n')
    for line in res:
        if len(line) > 0:
            if re.match(r"^tap.*", line):
                inf = re.sub(r"(tap[^:]+):", r"\1", line)
            elif re.match(r"^vnet.*", line):
                inf = re.sub(r"(vnet[^:]+):", r"\1", line)
            else:
                result[line] = inf

    return result
    
if __name__ == "__main__":
    if len(sys.argv) > 1:
        if sys.argv[1] in ['autoconf', 'detect']:
            if detect_kvm():
                print "yes"
            else:
                print "no"
        elif sys.argv[1] == "config":
            config(find_vm_names(list_pids()))
        else:
            fetch(find_vm_names(list_pids()))
    else:
        fetch(find_vm_names(list_pids()))