#!/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.
"""
import os
import sys
import warnings
from argparse import ArgumentParser
from getpass import getpass
from os import path
from proxmoxer import ProxmoxAPI

__author__ = 'Walter Doekes'
__copyright__ = 'Copyright (C) Walter Doekes, OSSO B.V. 2015-2016'
__licence__ = 'GPLv3+'
__version__ = '0.2.2'


# 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:
        value = row.get(item)
        if value is None:
            # Happens if the machine is new.
            pass
        elif value <= too_large:
            valid.append(value)
        else:
            ignored.append(value)
    return valid, ignored


def is_valid_uuid(uuid):
    parts = uuid.split('-')
    if ([len(i) for i in parts] == [8, 4, 4, 4, 12] and
            all(i in '0123456789abcdef-' for i in uuid)):
        return True
    return False


def get_rrddata(proxmox, timeframe='hour', aggregation='AVERAGE',
                only_vms=(), partial_match=False):
    vms = []
    crap = {}
    nodes_down = set()
    uuid_map = {}

    vms_data = proxmox.cluster.resources.get(type='vm')
    for i, vm in enumerate(vms_data):
        # Find only certain hosts?
        if only_vms:
            if (partial_match and
                    all(vm.get('name', '').find(i) == -1 for i in only_vms)):
                continue
            elif not partial_match and vm.get('name') not in only_vms:
                continue

        # VM info without status? Node down?
        if 'status' not in vm:
            if vm['node'] not in nodes_down:
                sys.stderr.write(
                    'WARNING: no status for vm %s. Is node %s down? '
                    'Ignoring further errors on that node.\n' % (
                        vm['vmid'], vm['node']))
                nodes_down.add(vm['node'])
            continue

        # Stopped VM?
        if vm['status'] in ('stopped',):
            continue

        assert vm['status'] in ('running', 'paused'), vm

        node = proxmox.nodes(vm['node'])
        container = getattr(node, vm['type'])(vm['vmid'])

        # Fetch UUIDs so we can complain about VMs without UUID.
        smbios1 = container.config.get().get('smbios1', '')
        uuid = ''.join(tmp[5:] for tmp in smbios1.split(',')
                       if tmp.startswith('uuid=')) or None
        assert vm['name'] not in uuid_map, (vm['name'], uuid_map)
        # We could choose to check manufacturer and product too:
        # uuid=648...a45,manufacturer=spindle,product=cacofonisk
        uuid_map[vm['name']] = uuid

        # # Update uuid if it wasn't set?
        # if not is_valid_uuid(uuid or ''):
        #     from uuid import uuid4 as make_uuid
        #     uuid = str(make_uuid())
        #     if smbios1:
        #         smbios1 = 'uuid=%s,%s' % (uuid, smbios1)
        #     else:
        #         smbios1 = 'uuid=%s' % (uuid,)
        #     try:
        #         container.config.post(smbios1=smbios1)
        #     except:
        #         sys.stderr.write(
        #             'Could not update smbios1 to %r on %s\n' % (
        #                 smbios1, vm['name']))
        #     else:
        #         print('UPDATED', smbios1, 'UPDATED on', vm['name'])

        # Get the RRD values.
        data = container.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)
            if not valid:
                sys.stderr.write(
                    'No sane RRD %s values for %s\n' % (item, vm['name']))
                valid = [0]
            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.
        if os.isatty(sys.stdout.fileno()):
            sys.stderr.write('% 3d%% loaded..\r' % (i * 100.0 / len(vms_data)))
            sys.stderr.flush()

    # Loop over uuid_map and complain about VMs without value UUID.
    valid_uuids, invalid_uuid_vms = [], []
    for name, uuid in uuid_map.items():
        if is_valid_uuid(uuid or ''):
            valid_uuids.append(uuid)
        else:
            invalid_uuid_vms.append(name)
    invalid_uuid_vms.sort()
    # Check validity.
    if invalid_uuid_vms:
        sys.stderr.write(
            'WARNING: These VMs have invalid/unset UUIDs: %s\n' % (
                ', '.join(invalid_uuid_vms),))
    # Check for dupes.
    if len(valid_uuids) != len(set(valid_uuids)):
        tmp = set()
        for test_uuid in valid_uuids:
            if test_uuid in tmp:
                tmp_vms = [name for name, uuid in uuid_map.items()
                           if uuid == test_uuid]
                sys.stderr.write(
                    'WARNING: Duplicate UUID %s in use by these VMs: %s\n' % (
                        test_uuid, ', '.join(tmp_vms)))
            tmp.add(test_uuid)

    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 subtype in ('avg', 'max'):
        for item, humanize in ITEMS:
            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[:443])')
    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(
        '--partial-match', default=False, action='store_true',
        help='Match VMs by substring instead of equality')
    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)

    # Extract optional port from hostname.
    if ':' in args.hostname:
        hostname, port = args.hostname.split(':', 1)
    else:
        hostname, port = args.hostname, 443  # used to be 8006

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

    # It's nice that urllib3 and requests show SNIMissingWarning and
    # InsecurePlatformWarning warnings, but we don't need them more than
    # once.
    warnings.simplefilter('once')

    # Log in, get the data, display it.
    api = ProxmoxAPI(hostname, port=port, 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, partial_match=args.partial_match)

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


if __name__ == '__main__':
    main()
