#!/usr/bin/env python
# vim: set ts=8 sw=4 sts=4 et ai:
from __future__ import print_function
"""
proxtop: Proxmox resource monitor -- list top resource consumers of your
Proxmox VM platform.

This is proxtop.  proxtop is free software: you can redistribute it
and/or modify it under the terms of the GNU General Public License as
published by the Free Software Foundation, version 3 or any later
version.
"""
__author__ = 'Walter Doekes'
__copyright__ = 'Copyright (C) Walter Doekes, OSSO B.V. 2015'
__licence__ = 'GPLv3+'
__version__ = '0.1.1'

import sys
from argparse import ArgumentParser
from getpass import getpass
from os import path
from proxmoxer import ProxmoxAPI

# Workaround for older python-requests with proxmoxer.
try:
    # python-requests 2.5.3
    from requests.packages import urllib3
except ImportError:
    # python-requests 2.2.1
    try:
        import requests
    except ImportError:
        # ProxmoxAPI will bail out later, saying:
        # Chosen backend requires 'requests' module
        pass
    else:
        import imp
        import urllib3

        class DummyException(Exception):
            pass
        dummy_func = (lambda *args, **kwargs: None)
        dummy_packages = imp.new_module('requests.packages')
        dummy_packages.urllib3 = urllib3
        urllib3.disable_warnings = dummy_func
        urllib3.exceptions.InsecureRequestWarning = DummyException
        requests.packages = dummy_packages


def humanize(multiplier, units):
    def wrapped(num):
        unitcopy = units[:]
        num *= multiplier
        while unitcopy:
            unit = unitcopy.pop()
            if num < 1024:
                break
            num /= 1024.0
        return '%6.1f %s' % (num, unit)
    return wrapped


def checking_foreach(data, item):
    """For some reason, the Proxmox API hands us crazy large values from
    time to time when we use the month-timeframe."""
    too_large = 0xffffffff  # 4 GiB/s..
    valid, ignored = [], []
    for row in data:
        if row[item] <= too_large:
            valid.append(row[item])
        else:
            ignored.append(row[item])
    return valid, ignored


def get_rrddata(proxmox, timeframe='hour', aggregation='AVERAGE',
                only_vms=()):
    vms = []
    crap = {}
    vms_data = proxmox.cluster.resources.get(type='vm')
    for i, vm in enumerate(vms_data):
        if only_vms and vm['name'] not in only_vms:
            continue
        if vm['status'] in ('stopped',):
            continue
        assert vm['status'] == 'running', vm

        rrddata = proxmox.nodes(vm['node']).openvz(vm['vmid']).rrddata
        data = rrddata.get(timeframe=timeframe, cf=aggregation)

        # Values tend to be wrong for only a single row: ignore the VM.
        if len(data) == 1:
            crap[vm['name']] = ['single_row', 1]
            vms.append((vm, dict((item, {'max': -1, 'avg': -1})
                                 for item, humanize in ITEMS)))
            continue

        # For 'hour' timeframe, we get 70 entries: one entry per minute:
        # len(data) == 70, (data[-1]['time'] - data[0]['time'] == 69 * 60),
        # For 'month' timeframe, we get ~66 entries: one entry per 12h, but
        # sometimes with gaps.
        totals = {}

        for item, humanize in ITEMS:
            valid, ignored = checking_foreach(data, item)
            totals[item] = {
                'max': max(valid),
                'avg': sum(valid) / float(len(valid)),
            }
            if ignored:
                if vm['name'] not in crap:
                    crap[vm['name']] = []
                for value in ignored:
                    crap[vm['name']].append((item, value))
        vms.append((vm, totals))

        # Progress.
        sys.stderr.write('% 3d%% loaded..\r' % (i * 100.0 / len(vms_data)))
        sys.stderr.flush()

    return vms, crap


def print_ignored(vms, crap, top):
    print('INGORED DATA: %d/%d' % (len(crap), len(vms)))
    print('------------------')
    crap = crap.items()
    crap.sort(key=(lambda x: len(x[1])), reverse=True)
    for i, (vm_name, values) in enumerate(crap[0:top]):
        print('#%d: %s  %r' %
              (i, vm_name, values))
    print()


def print_items(vms, top):
    for item, humanize in ITEMS:
        for subtype in ('avg', 'max'):
            print('SORTED BY: %s, %s' % (item, subtype))
            print('------------------')
            vms.sort(key=(lambda vm: vm[1][item][subtype]), reverse=True)
            for i, (vminfo, vmstat) in enumerate(vms[0:top]):
                print('#%d: %s  %s (%s)' %
                      (i, humanize(vmstat[item][subtype]),
                       vminfo['node'], vminfo['name']))
            print()


ITEMS = (
    ('cpu', (lambda x: '%6d %%    ' % (x * 100))),  # align with humanize
    ('diskread', humanize(1, ['TiB/s', 'GiB/s', 'MiB/s', 'KiB/s', 'B/s  '])),
    ('diskwrite', humanize(1, ['TiB/s', 'GiB/s', 'MiB/s', 'KiB/s', 'B/s  '])),
    ('netin', humanize(8, ['Tbps ', 'Gbps ', 'Mbps ', 'Kbps ', 'bps  '])),
    ('netout', humanize(8, ['Tbps ', 'Gbps ', 'Mbps ', 'Kbps ', 'bps  '])),
)


class ArgumentParser14191(ArgumentParser):
    """ArgumentParser from argparse that handles out-of-order positional
    arguments.

    This is a workaround created by Glenn Linderman in July 2012. You
    can now do this:

        parser = ArgumentParser14191()
        parser.add_argument('-f', '--foo')
        parser.add_argument('cmd')
        parser.add_argument('rest', nargs='*')
        # some of these would fail with the regular parser:
        for args, res in (('-f1 cmd 1 2 3', 'ok'),
                          ('cmd -f1 1 2 3', 'would_fail'),
                          ('cmd 1 -f1 2 3', 'would_fail'),
                          ('cmd 1 2 3 -f1', 'ok')):
            try: out = parser.parse_args(args.split())
            except: print 'args', 'failed', res
            # out: Namespace(cmd='cmd', foo='1', rest=['1', '2', '3'])

    Bugs: http://bugs.python.org/issue14191
    Files: http://bugs.python.org/file26273/t18a.py
    Changes: renamed to ArgumentParser14191 ** PEP cleaned ** hidden
      ErrorParser inside ArgumentParser14191 ** documented ** used
      new-style classes super calls  (Walter Doekes, March 2015)
    """
    class ErrorParser(ArgumentParser):
        def __init__(self, *args, **kwargs):
            self.__errorobj = None
            super(ArgumentParser14191.ErrorParser, self).__init__(
                *args, add_help=False, **kwargs)

        def error(self, message):
            if self.__errorobj:
                self.__errorobj.error(message)
            else:
                ArgumentParser.error(self, message)

        def seterror(self, errorobj):
            self.__errorobj = errorobj

    def __init__(self, *args, **kwargs):
        self.__setup = False
        self.__opt = ArgumentParser14191.ErrorParser(*args, **kwargs)
        super(ArgumentParser14191, self).__init__(*args, **kwargs)
        self.__opt.seterror(self)
        self.__setup = True

    def add_argument(self, *args, **kwargs):
        super(ArgumentParser14191, self).add_argument(*args, **kwargs)
        if self.__setup:
            chars = self.prefix_chars
            if args and len(args[0]) and args[0][0] in chars:
                self.__opt.add_argument(*args, **kwargs)

    def parse_args(self, args=None, namespace=None):
        ns, remain = self.__opt.parse_known_args(args, namespace)
        ns = super(ArgumentParser14191, self).parse_args(remain, ns)
        return ns


def main():
    # Fetch defaults from ~/.proxtoprc.
    defaults = {'hostname': None, 'username': None, 'password': None}
    proxtoprc = path.expanduser('~/.proxtoprc')
    if path.exists(proxtoprc):
        with open(proxtoprc, 'r') as file_:
            for line in file_:
                for key in defaults.keys():
                    if line.startswith(key + '='):
                        defaults[key] = line[(len(key) + 1):].strip()

    parser = ArgumentParser14191(
        description=('proxtop lists the top resource consumers on your '
                     'Proxmox VM platform.'),
        epilog=('Default values may be placed in ~/.proxtoprc. Lines should '
                'look like: hostname=HOSTNAME, username=USERNAME, '
                'password=PASSWORD'))
    parser.add_argument(
        'hostname', action='store',
        help='Use this API hostname (e.g. proxmox.example.com)')
    parser.add_argument(
        'username', action='store',
        help='Use this API username (e.g. monitor@pve)')
    parser.add_argument(
        '-T', '--top', default=8, action='store', type=int,
        help='Limit results to TOP VMs')
    parser.add_argument(
        '-t', '--timeframe', default='hour', action='store',
        help='Timeframe, can be one of: hour* | day | week | month | year')
    parser.add_argument(
        '-g', '--aggregation', default='AVERAGE', action='store',
        help='RRD aggregation, can be one of: AVERAGE* | MAX')
    parser.add_argument(
        'only_vms', nargs='*',
        help='Limit results to these VMs')

    # Pass adjusted argv, based on defaults found in ~/.proxtoprc.
    argv = sys.argv[1:]
    if defaults['hostname'] and defaults['username']:
        argv.insert(0, defaults['hostname'])
        argv.insert(1, defaults['username'])
    args = parser.parse_args(argv)

    # Fetch password.
    if defaults['password']:
        args.password = defaults['password']
    else:
        args.password = getpass('Password:')

    # Log in, get the data, display it.
    api = ProxmoxAPI(args.hostname, user=args.username,
                     password=args.password, verify_ssl=False)

    vms, crap = get_rrddata(
        api, timeframe=args.timeframe, aggregation=args.aggregation,
        only_vms=args.only_vms)

    if crap:
        print_ignored(vms, crap, top=args.top)
    print_items(vms, top=args.top)


if __name__ == '__main__':
    main()
