#!/usr/bin/python3

#Import librarys
import inspect
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))))

import threading
from http.server import HTTPServer
import socket
import time

from prometheus.collectors import Gauge
from prometheus.registry import Registry
from prometheus.exporter import PrometheusMetricHandler

import requests
import json
import configparser

# Parse and read config file
config = configparser.ConfigParser()
config.read(os.path.expanduser("/root/prometheus/config.conf"))
server_name = config.get("login", "server")
node = config.get("login", "node")
username = config.get("login", "username")
password = config.get("login", "password")
PORT_NUMBER = int(config.get("server", "port"))

#set start time for geting new key after 2 hours
start_time = time.time()

# api-endpoint
url = server_name + "/api2/json/access/ticket"

#PARAMS = {'data':auth}
data = {'username':username,
        'password': password}

#ticket = response['data']['ticket']
#token = response['data']['CSRFPreventionToken']


#get API security code (ticket)
def get_key():
    global start_time
    start_time = time.time()
    r = requests.post(url = url, data = data)
    response = r.json()
    ticket = response['data']['ticket']
    return ticket

#set cookies for authorization with ticket code
cookies = dict(PVEAuthCookie= get_key())

#refresh key and cookie variable
def refresh_key():
    global cookies
    cookies = dict(PVEAuthCookie= get_key())


# get request function
def get_request(api_url):
    r = requests.get(server_name + api_url, cookies=cookies)
    return r.json()

# get node status
def get_node_status():
    node_status = get_request("/api2/json/nodes/" + node + "/status")
    proxmox_host_metric = node_status['data']
    return proxmox_host_metric

# get all VM metric
def get_all_vm():
    all_vm = get_request("/api2/json/cluster/resources?type=vm")
    vm = all_vm['data']
    list = [item for item in vm if item['status'] == "running"]
    return list


# Gether metrics
def gather_data(registry):


    # Get the host name of the machine
    host = socket.gethostname()

    # Create our collectors
    ram_metric = Gauge("memory_usage_bytes", "Memory usage in bytes",
                       {'host': host})
    max_memory_metric = Gauge("max_memory_assigned", "Assigned memmory in bytes",
                       {'host': host})
    cpu_metric = Gauge("cpu_usage_percent", "CPU usage percent.",
                       {'host': host})
    disk_read_metric = Gauge("disk_read_bytes", "Disk read in bytes",
                       {'host': host})
    disk_write_metric = Gauge("disk_write_bytes", "Disk write in bytes",
                       {'host': host})
    net_in_metric = Gauge("net_in_bytes", "Network input in bytes",
                       {'host': host})
    net_out_metric = Gauge("net_out_bytes", "Network output in bytes",
                       {'host': host})
    # pmx host collectors
    pmx_root_fs = Gauge("pmx_root_fs_bytes", "Root File System in bytes",
                       {'host': host})
    pmx_load_avg = Gauge("pmx_load_avg", "Proxmox load average",
                       {'host': host})
    pmx_cpu_usage  = Gauge("pmx_cpu_usage_percentage", "Proxmox cpu usage in percentage",
                       {'host': host})
    pmx_memory_usage = Gauge("pmx_memory_usage_bytes", "Proxmox memory usage in bytes",
                       {'host': host})


    # register the metric collectors
    registry.register(ram_metric)
    registry.register(max_memory_metric)
    registry.register(cpu_metric)
    registry.register(disk_read_metric)
    registry.register(disk_write_metric)
    registry.register(net_in_metric)
    registry.register(net_out_metric)
    # register proxmox metrics
    registry.register(pmx_root_fs)
    registry.register(pmx_load_avg)
    registry.register(pmx_cpu_usage)
    registry.register(pmx_memory_usage)


    # Start gathering metrics every 15 seconds
    while True:
        if time.time() - start_time >= 3600:
            refresh_key()

        time.sleep(15)

        # Add metrics
        metric = get_all_vm()
        proxmox_metric = get_node_status()

        for item in metric:

            for key,value in item.items():
                if key == "mem":
                    ram_metric.set({'name': item['name'], 'type': 'usage'}, item['mem'])
                elif key == "maxmem":
                    max_memory_metric.set({'name': item['name'], 'type': 'max'}, item['maxmem'])
                elif key == "cpu":
                    cpu_metric.set({'name': item['name'], 'type': 'load'}, item['cpu'])
                elif key == "diskread":
                    disk_read_metric.set({'name': item['name'], 'type': 'disk_read'}, item['diskread'])
                elif key == "diskwrite":
                    disk_write_metric.set({'name': item['name'], 'type': 'disk_write'}, item['diskwrite'])
                elif key == "netin":
                    net_in_metric.set({'name': item['name'], 'type': 'net_in'}, item['netin'])
                elif key == "netout":
                    net_out_metric.set({'name': item['name'], 'type': 'net_out'}, item['netout'])

        pmx_root_fs.set({'name': host, 'type': 'Free'}, proxmox_metric['rootfs']['free'])
        pmx_root_fs.set({'name': host, 'type': 'Total'}, proxmox_metric['rootfs']['total'])
        pmx_root_fs.set({'name': host, 'type': 'Used'}, proxmox_metric['rootfs']['used'])
        pmx_root_fs.set({'name': host, 'type': 'Available'}, proxmox_metric['rootfs']['avail'])

        pmx_load_avg.set({'name': host, 'type': '1 minute'}, float(proxmox_metric['loadavg'][0]))
        pmx_load_avg.set({'name': host, 'type': '5 minutes'}, float(proxmox_metric['loadavg'][1]))
        pmx_load_avg.set({'name': host, 'type': '15 minutes'}, float(proxmox_metric['loadavg'][2]))

        pmx_cpu_usage.set({'name': host, 'type': 'idle'}, proxmox_metric['cpu'])
        pmx_cpu_usage.set({'name': host, 'type': 'IO wait'}, proxmox_metric['wait'])

        pmx_memory_usage.set({'name': host, 'type': 'used'}, proxmox_metric['memory']['used'])
        pmx_memory_usage.set({'name': host, 'type': 'free'}, proxmox_metric['memory']['free'])
        pmx_memory_usage.set({'name': host, 'type': 'total'}, proxmox_metric['memory']['total'])


if __name__ == "__main__":

    # Create the registry
    registry = Registry()

    # Create the thread that gathers the data while we serve it
    thread = threading.Thread(target=gather_data, args=(registry, ))
    thread.start()

    # Set a server to export (expose to prometheus) the data (in a thread)
    try:
        # We make this to set the registry in the handler
        def handler(*args, **kwargs):
            PrometheusMetricHandler(registry, *args, **kwargs)

        server = HTTPServer(('', PORT_NUMBER), handler)
        server.serve_forever()

    except KeyboardInterrupt:
        server.socket.close()
        thread.join()