#!/usr/bin/env python3

import argparse
import os
import sys
import subprocess
import json
import hashlib


def perform_system_checks():
    if sys.version_info <= (3, 7):
        print("Tool needs to be run with at least Python 3.7")
        sys.exit(1)

    version_proc = subprocess.run(['which', 'VBoxManage'], capture_output=True)
    if not version_proc.returncode == 0:
        print("VBoxManage must be installed")
        sys.exit(1)
    qemuimg_proc = subprocess.run(['which', 'qemu-img'],  capture_output=True)
    if not qemuimg_proc.returncode == 0:
        print("qemu-img must be installed for this program to work.")
        sys.exit(1)


def parse_arguments():
    parser = argparse.ArgumentParser(
        description=(
            "Given a virtual disk image file and a name for the VM, "
            "produce an OVA file ready for deployment in VMware ESXi "
            "and other platforms."
        )
    )
    parser.add_argument(
        "--diskfile",
        "-d",
        type=str,
        help=(
            "File name of the virtual disk image. OVA file will be created "
            "in the same directory. This defaults to '<name>.vmdk'. If "
            "the extension is something other than 'vmdk', the file will be "
            "converted before packing and the files will be placed into "
            "the current directory."
        ),
    )
    parser.add_argument(
        "--template",
        "-t",
        type=str,
        help=(
            "File name of the XML template file used for OVA creation."
        ),
        default="config/ova-template.xml",
    )
    parser.add_argument(
        "--vm-name",
        "-n",
        type=str,
        help="Name of the VM (should be the pfsystemname)",
        required=True
    )
    parser.add_argument(
        "--num-cpus",
        "-c",
        type=int,
        help="Number of CPU cores in the VM",
        default=4,
    )
    parser.add_argument(
        "--ram-mb",
        "-r",
        type=int,
        help="Amount of RAM in the VM in MB",
        default=16384,
    )
    parser.add_argument(
        "--upload",
        "-u",
        help="upload finished OVA to OwnCloud",
        action="store_true"
    )
    parser.add_argument(
        "--force",
        "-f",
        help="force overwriting an existing OVA file",
        action="store_true"
    )
    args = parser.parse_args()
    if args.diskfile is None:
        args.diskfile = args.name + '.vmdk'
    return args


def make_ova(abs_diskfile, output_dir, ova_file, vm_name,
             num_cpus=2, ram_mb=8192):

    if not os.path.isfile(abs_diskfile):
        print("Disk file not found: {}".format(abs_diskfile))
        sys.exit(1)

    ovf_file = vm_name + '.ovf'
    vmdk_file = vm_name + '.vmdk'
    vmdk_stream_file = vm_name + '-disk1.vmdk'
    manifest_file = vm_name + '.mf'

    if abs_diskfile != vmdk_file:
        if not abs_diskfile.endswith('.vmdk'):
            subprocess.run([
                'qemu-img', 'convert', abs_diskfile, '-p', '-O', 'vmdk',
                vmdk_file
            ], check=True, cwd=output_dir)
        else:
            os.link(
                os.path.join(output_dir, abs_diskfile),
                os.path.join(output_dir, vmdk_file)
            )

    if os.path.isfile(os.path.join(output_dir, ova_file)):
        os.remove(os.path.join(output_dir, ova_file))

    vmdk_info_proc = subprocess.run(
        ['qemu-img', 'info', vmdk_file, '--output=json'],
        check=True, cwd=output_dir, capture_output=True, text=True,
    )
    vmdk_info = json.loads(vmdk_info_proc.stdout)

    # Generate streaming mode vmdk file
    subprocess.run(
        [
            'VBoxManage', 'clonemedium', vmdk_file, vmdk_stream_file,
            '--format', 'VMDK', '--variant', 'Stream'
        ],
        check=True, cwd=output_dir,
    )
    # VBoxManage keeps a handle on the files it touches, which musst be
    # explicitly removed.
    for close_file in (vmdk_file, vmdk_stream_file):
        subprocess.run(
            ['VBoxManage', 'closemedium', close_file],
            check=True, cwd=output_dir,
        )
    os.remove(os.path.join(output_dir, vmdk_file))

    vmdk_size = os.stat(os.path.join(output_dir, vmdk_stream_file)).st_size

    template = open(args.template, 'r').read()
    ovf_data = template.format(
        vmdk_file=vmdk_stream_file,
        vmdk_size=vmdk_size,
        disk_size=vmdk_info['virtual-size'],
        populated_size=vmdk_info['actual-size'],
        vm_name=vm_name,
        num_cpus=num_cpus,
        ram_mb=ram_mb,
    )
    with open(os.path.join(output_dir, ovf_file), 'w') as fd:
        fd.write(ovf_data.strip())

    # Write the manifest file
    checksums = ''
    for filename in (ovf_file, vmdk_stream_file):
        sha256 = hashlib.sha256()
        with open(os.path.join(output_dir, filename), 'rb') as fd:
            for block in iter(lambda: fd.read(65536), b''):
                sha256.update(block)
        hash = sha256.hexdigest()
        checksums += f'SHA256({filename})= {hash}\n'

    with open(os.path.join(output_dir, manifest_file), 'w') as manifest_fh:
        manifest_fh.write(checksums)

    # Make the tar archive
    subprocess.run(
        [
            'tar', 'cf', ova_file,
            '--owner=someone', '--group=someone',
            ovf_file, manifest_file, vmdk_stream_file,
        ],
        check=True, cwd=output_dir,
    )

    for cleanup in (ovf_file, vmdk_stream_file, manifest_file):
        os.remove(os.path.join(output_dir, cleanup))


def upload(ova_file, output_dir):
    subprocess.run([
        'rsync', '--progress', '--partial', ova_file,
        'perfact@tyke:/vol/nextcloud/data/pfsync/files/'
    ], cwd=output_dir, check=True)
    subprocess.run([
        'ssh', 'perfact@tyke', 'sudo', '-u', 'www-data',
        'php', '/vol/nextcloud/occ', 'files:scan', 'pfsync'
    ], check=True)


if __name__ == '__main__':
    perform_system_checks()
    args = parse_arguments()

    abs_diskfile = os.path.abspath(args.diskfile)
    if abs_diskfile.endswith('.vmdk'):
        output_dir = os.path.dirname(abs_diskfile)
    else:
        output_dir = os.path.abspath('.')

    ova_file = args.vm_name + '.ova'

    if not os.path.isfile(os.path.join(output_dir, ova_file)) or args.force:
        make_ova(
            abs_diskfile=abs_diskfile,
            output_dir=output_dir,
            ova_file=ova_file,
            vm_name=args.vm_name,
            num_cpus=args.num_cpus,
            ram_mb=args.ram_mb,
        )
    else:
        if args.upload:
            print("Uploading old OVA file, use --force to overwrite.")
        else:
            print("OVA file is already present, use --force to overwrite.")

    if args.upload:
        upload(ova_file=ova_file, output_dir=output_dir)