#!/usr/bin/env python3
# Let's stick to core python3 modules
import argparse
import gettext
import hashlib
import http.client
import io
import json
import os
import shutil
import socket
from subprocess import check_output
import sys
import tarfile
import tempfile
import urllib.request
import uuid


DEFAULT_VGNAME = "LXDStorage"

_ = gettext.gettext
gettext.textdomain("lxd")


class FriendlyParser(argparse.ArgumentParser):
    def error(self, message):
        sys.stderr.write('error: %s\n' % message)
        self.print_help()
        sys.exit(2)


class UnixHTTPConnection(http.client.HTTPConnection):
    def __init__(self, path):
        http.client.HTTPConnection.__init__(self, 'localhost')
        self.path = path

    def connect(self):
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        sock.connect(self.path)
        self.sock = sock

        
class LXD(object):
    def __init__(self, path):
        self.lxd = UnixHTTPConnection(path)

    def rest_call(self, path, data=None, method="GET", headers={}):
        if method == "GET" and data:
            self.lxd.request(
                method,
                "%s?%s" % "&".join(["%s=%s" % (key, value)
                                    for key, value in data.items()]), headers)
        else:
            self.lxd.request(method, path, data, headers)

        r = self.lxd.getresponse()
        d = json.loads(r.read().decode("utf-8"))
        return r.status, d

    def set_lvm_vgname(self, vgname):
        self._set_lvm_config("storage.lvm_vg_name", vgname)

    def set_lvm_poolname(self, poolname):
        self._set_lvm_config("storage.lvm_thinpool_name", poolname)

    def _set_lvm_config(self, key, val):
        data = json.dumps({"config": {key: val}})

        status, data = self.rest_call("/1.0", data, "PUT")

        if status != 200:
            sys.stderr.write("Error in setting vgname:{}\n{}\n".format(status,
                                                                   data))
            raise Exception("Failed to set vgname: %s" % vgname)

    def get_server_config(self):
        status, config = self.rest_call("/1.0", "", "GET")
        if status != 200:
            sys.stderr.write("Error in setting vgname:{}\n{}\n".format(status,
                                                                   data))
            raise Exception("Failed to set vgname: %s" % vgname)

        return config["metadata"]["config"]


def lxd_dir():
    if "LXD_DIR" in os.environ:
        return os.environ["LXD_DIR"]
    else:
        return "/var/lib/lxd"


def connect_to_socket():
    lxd_socket = os.path.join(lxd_dir(), "unix.socket")

    if not os.path.exists(lxd_socket):
        print(_("LXD isn't running."))
        sys.exit(1)

    return LXD(lxd_socket) 
    
    
def create_image(args):
    imgfname = os.path.join(lxd_dir(), "{}.img".format(args.size))
    rollbacks = []
    try:
        print("Creating sparse backing file {}".format(imgfname), flush=True)
        check_output("truncate -s {} {}".format(args.size, imgfname), shell=True)
        rollbacks.append("rm {}".format(imgfname))

        print("Setting up loop device", flush=True)
        pvloopdev = check_output("losetup -f", shell=True).decode().strip()
        check_output("losetup {} {}".format(pvloopdev, imgfname), shell=True)
        rollbacks.append("losetup -d " + pvloopdev)

        print("Creating LVM PV {}".format(pvloopdev), flush=True)
        check_output("pvcreate {}".format(pvloopdev), shell=True)
        rollbacks.append("pvremove " + pvloopdev)

        print("Creating LVM VG {}".format(DEFAULT_VGNAME), flush=True)
        check_output("vgcreate {} {}".format(DEFAULT_VGNAME, pvloopdev), shell=True)
        rollbacks.append("vgremove {}".format(DEFAULT_VGNAME))

    except Exception as e:
        sys.stderr.write("Error: {}. Cleaning up:\n".format(e))
        for rbcmd in reversed(rollbacks):
            sys.stderr.write("+ {}\n".format(rbcmd))
            check_output(rbcmd, shell=True)
        raise e

def destroy_image(args, lxd):
    print("Checking current LXD configuration", flush=True)
    cfg = lxd.get_server_config()
    vgname = cfg.get("storage.lvm_vg_name", None)
    if vgname is None:
        sys.stderr.write("LXD is not configured for LVM. No changes will be made.\n")
        return

    lvnames = check_output("lvs {} -o name,lv_attr --noheadings".format(vgname), shell=True).decode().strip()
    used_lvs = []
    for lvline in lvnames.split("\n"):
        if lvline == '':
            continue
        name, attrs = lvline.split()
        if attrs.strip().startswith("V"):
            used_lvs.append(name)
    if len(used_lvs) > 0:
        print("LVM storage is still in use by the following volumes: {}".format(used_lvs))
        print("Please delete the corresponding images and/or containers before destroying storage.")
        sys.exit()

    pvname = check_output("sudo vgs {} --noheadings -o pv_name".format(vgname), shell=True).decode().strip()
    print("Removing volume group {}".format(vgname))
    check_output("vgremove -f {}".format(vgname), shell=True)
    print("Removing physical volume {}".format(pvname))
    check_output("pvremove -y {}".format(pvname), shell=True)
    
    imgfname = check_output("losetup {} --list --output BACK-FILE --noheadings".format(pvname), shell=True).decode().strip()
    print("Detaching loop device {}".format(pvname))
    check_output("losetup -d {}".format(pvname), shell=True)
    print("Deleting backing file {}".format(imgfname))
    check_output("rm {}".format(imgfname), shell=True)

        
def do_main():
    
    parser = FriendlyParser(
        description=_("LXD: LVM storage helper"),
        formatter_class=argparse.RawTextHelpFormatter,
        epilog=_("""Examples:
 To create a 10G sparse loopback file and register it with LVM and LXD:
    %s -s 10G
 To de-configure LXD and destroy the LVM volumes and backing file:
    %s --destroy
""" % (sys.argv[0], sys.argv[0])))
    parser.add_argument("-s", "--size", help=_("Size of backing file to register as LVM PV"))
    parser.add_argument("--destroy", action="store_true", default=False, help=_("Un-configure LXD and delete image file"))

    args = parser.parse_args()
    if os.geteuid() != 0:
        sys.exit("Configuring LVM requires root privileges.")

    try:
        check_output("type vgcreate", shell=True)
    except:
        sys.exit("lvm2 tools not found. try 'apt-get install lvm2'")

    lxd = connect_to_socket()

    if args.destroy:
        try:
            destroy_image(args, lxd)
            print("Clearing LXD storage configuration")
            lxd.set_lvm_vgname("")
            lxd.set_lvm_poolname("")
        except Exception as e:
            sys.stderr.write("Error destroying image:")
            sys.stderr.write(str(e))
            sys.stderr.write("\n")

    else:
        try:
            create_image(args)
        except:
            sys.stderr.write("Stopping.\n")
        else:
            try:
                print("Configuring LXD")
                lxd.set_lvm_vgname(DEFAULT_VGNAME)
            except:
                sys.stderr.write("Error configuring LXD, removing backing file\n")
                destroy_image(args, lxd)

    print("Done.")


if __name__ == "__main__":
    do_main()

