#!/usr/bin/env python3

# (c) 2020 Heinlein Support GmbH
#          Robert Sander <r.sander@heinlein-support.de>

# This 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 in version 2.  This file is distributed
# in the hope that it will be useful, but WITHOUT ANY WARRANTY;  with-
# out even the implied warranty of  MERCHANTABILITY  or  FITNESS FOR A
# PARTICULAR PURPOSE. See the  GNU General Public License for more de-
# ails.  You should have  received  a copy of the  GNU  General Public
# License along with GNU Make; see the file  COPYING.  If  not,  write
# to the Free Software Foundation, Inc., 51 Franklin St,  Fifth Floor,
# Boston, MA 02110-1301 USA.

from shutil import which
from psutil import disk_partitions, disk_usage
from os import walk
from os.path import join, isfile
from socket import gethostname, getfqdn
import subprocess
import json

from pprint import pprint

def parse_output(cmd, parse_json=True):
    r = subprocess.run(cmd, capture_output=True, check=True, text=True)
    if parse_json:
        return json.loads(r.stdout)
    else:
        return r.stdout

def get_filesystem_type(partitions, file):
    for p in partitions:
        if file.startswith(p.mountpoint):
            return p.fstype
    return None

def df(partitions, file):
    capacity, used, free, percent = disk_usage(file)

    typ = get_filesystem_type(partitions, file)

    return typ, capacity, used

def du(directory):
    du = which("du")
    if not du:
        raise RuntimeError("du not found")
    size, _ = parse_output([du, "-B", "1", "-s", directory], False).split()
    return int(size)

def provisioned_space(directory):
    prov = 0
    qemuimg = which("qemu-img")
    if not qemuimg:
        return prov
    for root, dirs, files in walk(directory):
        for file in files:
            path = join(root, file)
            if isfile(path):
                for line in parse_output([qemuimg, "info", path], False).split("\n"):
                    if line.startswith("virtual size: "):
                        prov += int(line.split("(")[1].split(" ")[0])
    return prov

def handle_dir_storage(partitions, path):
    images = path + "/images"
    typ, capacity, total_used = df(partitions, images)
    used = du(images)
    prov = provisioned_space(images)
    total_cap = capacity - total_used + used
    return typ, total_cap, used, prov

def handle_zfs_storage(pool):
    zfs = which("zfs")
    if not zfs:
        raise RuntimeError("zfs not found")
    free = int(parse_output([zfs, "get", "-Hpo", "value", "available", pool], False))
    used = int(parse_output([zfs, "get", "-Hpo", "value", "used", pool], False))
    prov = sum(map(int, parse_output([zfs, "get", "-rHpo", "value", "volsize", pool], False).split("\n")[1:-1]))
    return "ZFS", free + used, used, prov

if __name__ == "__main__":
    pvesh = which("pvesh")
    if pvesh:
        partitions = sorted(disk_partitions(True), key=lambda x: x.mountpoint, reverse=True)
        hostname = gethostname().split(".", 1)[0]
        fqdn = getfqdn()
        print("<<<esx_vsphere_datastores>>>")
        for storage_info in parse_output([pvesh, "get", "/storage", "--output-format=json"]):
            nodes = storage_info.get("nodes", "").split(",")
            if "images" in storage_info["content"] and ( hostname in nodes or fqdn in nodes):
                typ, capacity, used, provisioned = None, 0, 0, 0
                if storage_info["type"] in ["dir", "nfs", "cephfs"]:
                    typ, capacity, used, provisioned = handle_dir_storage(partitions, storage_info["path"])
                    url = storage_info["path"]
                if storage_info["type"] == "zfspool":
                    typ, capacity, used, provisioned = handle_zfs_storage(storage_info["pool"])
                    url = storage_info["pool"]
                if typ:
                    print("[%s]" % storage_info["storage"])
                    print("url %s" % url)
                    print("accessible %s" % (storage_info.get("disable", 0) == 0) )
                    print("type %s" % typ)
                    print("capacity %d" % capacity)
                    print("freeSpace %d" % (capacity - used))
                    if provisioned > 0:
                        print("uncommitted %d" % (provisioned - used))
                    else:
                        print("uncommitted 0")
    pvesm = which("pvesm")
    if pvesm:
        print("<<<esx_vsphere_datastores>>>")
        for line in parse_output([pvesm, "status"], False).split("\n")[1:-1]:
            dsname, typ, status, total_cap, used, avail, perc = line.split()
            if typ == "lvm":
                print(f"[{dsname}]")
                print("accessible %s" % (status == "active") )
                print("type %s" % typ)
                print("capacity %d" % (int(total_cap) * 1024))
                print("freeSpace %d" % (int(avail) * 1024))
                print("uncommitted 0")

