#!/usr/bin/env python3

# See docs/vmware-ova.md for info

import argparse
import os
import sys
import subprocess
import json
import hashlib
import tempfile
import tarfile

try:
    from mako.template import Template
except ImportError:
    sys.exit("Mako Templates module not found")

MIN_PYTHON = (3, 7)
GUEST_OS = {
    "debian10_64Guest": {
        "definition": "96",
        "description": "Garden Linux 64 (Debian 11)",
        "hostname": "gardenlinux"
    },
    "sles15_64Guest": {
        "definition": "85",
        "description": "SUSE Linux Enterprise Server 15 SP2 (CHOST)",
        "hostname": "sles15chost"
    }
}
PROPERTIES = {
    "vmdkFilename": "",
    "vmdkFileSize": "",
    "vmdkFileCapacity": "",
    "virtualSystemId": "",
    "osType": "",
    "osDefinition": "",
    "osDescription": "",
    "applianceHostname": ""
}


class OVAImageBuild:
    def getInfo(self, filename):
        result = subprocess.run(["qemu-img", "info", "--output", "json", filename], capture_output=True)
        if result.returncode != 0:
            sys.exit("Error getting image info for " + filename + ": " + result.stdout)

        doc = json.loads(result.stdout)
        fileSize = os.stat(filename).st_size
        virtualSize = doc["format-specific"]["data"]["extents"][0]["virtual-size"]
        return (fileSize, virtualSize)

    def makeAttributes(self, PROPERTIES):
        class P:
            pass
        p = P()
        for k in PROPERTIES:
            setattr(p, k, PROPERTIES[k])
        return p

    def sha256(self, filename):
        blocksize = 65536
        sha = hashlib.sha256()
        with open(filename, 'rb') as file:
            fileBuffer = file.read(blocksize)
            while len(fileBuffer) > 0:
                sha.update(fileBuffer)
                fileBuffer = file.read(blocksize)

        return sha.hexdigest()

    def run(self, template):
        vmdkFullFilename = PROPERTIES["vmdkFullFilename"]
        vmdkPath = os.path.dirname(vmdkFullFilename)
        vmdkFilename = os.path.basename(PROPERTIES["vmdkFilename"])
        ovaFilename = os.path.splitext(vmdkFilename)[0] + ".ova"
        ovaFullFilename = os.path.join(vmdkPath, ovaFilename)
        ovfFilename = os.path.splitext(vmdkFilename)[0] + ".ovf"
        print(ovaFullFilename)
        fileSize, virtualSize = self.getInfo(PROPERTIES["vmdkFullFilename"])
        PROPERTIES["vmdkFileSize"] = fileSize
        PROPERTIES["vmdkFileCapacity"] = virtualSize
        p = self.makeAttributes(PROPERTIES)
        tmpl = Template(filename=template)
        ovf = tmpl.render(p=p)
        vmdkSha256 = self.sha256(vmdkFullFilename)
        with tempfile.TemporaryDirectory() as tmpDir:

            ovfFullFilename = os.path.join(tmpDir, ovfFilename)
            with open(ovfFullFilename, "w") as file:
                file.write(ovf)
            ovfSha256 = self.sha256(ovfFullFilename)

            mfFilename = os.path.splitext(vmdkFilename)[0] + ".mf"
            mfFullFilename = os.path.join(tmpDir, mfFilename)

            mf = "SHA256(" + vmdkFilename + ")= " + vmdkSha256 + "\n" + "SHA256(" + ovfFilename + ")= " + ovfSha256 + "\n"

            with open(mfFullFilename, "w") as file:
                file.write(mf)

            with tarfile.open(ovaFullFilename, "w") as tar:
                tar.add(ovfFullFilename, arcname=ovfFilename)
                tar.add(vmdkFullFilename, arcname=vmdkFilename)
                tar.add(mfFullFilename, arcname=mfFilename)

    @classmethod
    def _argparse_register(cls, parser):
        parser.add_argument(
            '--vmdk',
            type=str,
            dest='vmdk',
            help='vmdk filename',
            required=True
        )
        parser.add_argument(
            "--guest-id",
            type=str,
            dest='guest_id',
            help='Guest operating system identifier',
            choices=list(GUEST_OS.keys()),
            required=True
        )
        parser.add_argument(
            "--template",
            type=str,
            dest='template',
            help='OVF template file path (see ../templates for examples)',
            required=True
        )

    @classmethod
    def _main(cls):
        parser = argparse.ArgumentParser()
        cls._argparse_register(parser)
        args = parser.parse_args()

        if sys.version_info < MIN_PYTHON:
            sys.exit("Python %s.%s or newer is required." % MIN_PYTHON)

        PROPERTIES["vmdkFullFilename"] = args.vmdk
        PROPERTIES["vmdkFilename"] = os.path.basename(args.vmdk)
        PROPERTIES["virtualSystemId"] = os.path.splitext(os.path.basename(args.vmdk))[0]
        PROPERTIES["osType"] = args.guest_id
        PROPERTIES["osDefinition"] = GUEST_OS[args.guest_id]["definition"]
        PROPERTIES["osDescription"] = GUEST_OS[args.guest_id]["description"]
        PROPERTIES["applianceHostname"] = GUEST_OS[args.guest_id]["hostname"]
        oVAImageBuild = cls()
        oVAImageBuild.run(template=args.template)


if __name__ == '__main__':
    OVAImageBuild._main()
