#!/usr/bin/python3

 
 import os
 import re
 import shutil
 import tarfile
 import fileinput

 from argparse import ArgumentParser

 import json

 from utils import Utils

 
 def create_ova_image(raw_image_name, tools_path, build_scripts_path, config):
     output_path = os.path.dirname(os.path.realpath(raw_image_name))
     utils = Utils()
     # Remove older artifacts

     files = os.listdir(output_path)

     for file in files:

         if file.endswith(".vmdk"):

             os.remove(os.path.join(output_path, file))
 
     vmx_path = output_path + '/photon-ova.vmx'

     utils.replaceandsaveasnewfile(build_scripts_path + '/vmx-template',
                                   vmx_path, 'VMDK_IMAGE',
                                   output_path + '/photon-ova.vmdk')

     vixdiskutil_path = tools_path + 'vixdiskutil'
     vmdk_path = output_path + '/photon-ova.vmdk'
     ovf_path = output_path + '/photon-ova.ovf'
     mf_path = output_path + '/photon-ova.mf'

     ovfinfo_path = build_scripts_path + '/ovfinfo.txt'

     vmdk_capacity = (int(config['size']['root']) +
                      int(config['size']['swap'])) * 1024
     utils.runshellcommand(
         "{} -convert {} -cap {} {}".format(vixdiskutil_path,
                                            raw_image_name,
                                            vmdk_capacity,
                                            vmdk_path))
     utils.runshellcommand(
         "{} -wmeta toolsVersion 2147483647 {}".format(vixdiskutil_path, vmdk_path))

 
     utils.runshellcommand("ovftool {} {}".format(vmx_path, ovf_path))
     utils.replaceinfile(ovf_path, 'otherGuest', 'other3xLinux64Guest')
 
     #Add product info

     if os.path.exists(ovfinfo_path):
         with open(ovfinfo_path) as f:
             lines = f.readlines()
             for line in fileinput.input(ovf_path, inplace=True):

                 if line.strip() == '</VirtualHardwareSection>':

                     for ovfinfoline in lines:

                         print(ovfinfoline)

                 else:

                     print(line)

 
     if os.path.exists(mf_path):
         os.remove(mf_path)
 
     cwd = os.getcwd()
     os.chdir(output_path)
     out = utils.runshellcommand("openssl sha1 photon-ova-disk1.vmdk photon-ova.ovf")
     with open(mf_path, "w") as source:
         source.write(out)
     rawsplit = os.path.splitext(raw_image_name)
     ova_name = rawsplit[0] + '.ova'
 

     ovatar = tarfile.open(ova_name, "w", format=tarfile.USTAR_FORMAT)

     for name in ["photon-ova.ovf", "photon-ova.mf", "photon-ova-disk1.vmdk"]:
         ovatar.add(name, arcname=os.path.basename(name))
     ovatar.close()

     os.remove(vmx_path)

     os.remove(mf_path)
 
     if 'additionalhwversion' in config:
         for addlversion in config['additionalhwversion']:

             new_ovf_path = output_path + "/photon-ova-hw{}.ovf".format(addlversion)
             mf_path = output_path + "/photon-ova-hw{}.mf".format(addlversion)

             utils.replaceandsaveasnewfile(
                 ovf_path, new_ovf_path, "vmx-.*<", "vmx-{}<".format(addlversion))
             out = utils.runshellcommand("openssl sha1 photon-ova-disk1.vmdk "
                                         "photon-ova-hw{}.ovf".format(addlversion))

             with open(mf_path, "w") as source:
                 source.write(out)
             temp_name_list = os.path.basename(ova_name).split('-')
             temp_name_list = temp_name_list[:2] + ["hw{}".format(addlversion)] + temp_name_list[2:]
             new_ova_name = '-'.join(temp_name_list)

             new_ova_path = output_path + '/' + new_ova_name

             ovatar = tarfile.open(new_ova_path, "w", format=tarfile.USTAR_FORMAT)

             for name in [new_ovf_path, mf_path, "photon-ova-disk1.vmdk"]:
                 ovatar.add(name, arcname=os.path.basename(name))
             ovatar.close()
 
             os.remove(new_ovf_path)
             os.remove(mf_path)

     os.chdir(cwd)

     os.remove(ovf_path)
     os.remove(vmdk_path)

     files = os.listdir(output_path)

     for file in files:
         if file.endswith(".vmdk"):

             os.remove(os.path.join(output_path, file))

 
 
 if __name__ == '__main__':
     usage = "Usage: %prog [options]"

     parser = ArgumentParser(usage)
 
     parser.add_argument("-r", "--raw-image-path", dest="raw_image_path")
     parser.add_argument("-c", "--vmdk-config-path", dest="vmdk_config_path")
     parser.add_argument("-w", "--working-directory", dest="working_directory")
     parser.add_argument("-m", "--mount-path", dest="mount_path")
     parser.add_argument("-a", "--additional-rpms-path", dest="additional_rpms_path")
     parser.add_argument("-i", "--image-name", dest="image_name")
     parser.add_argument("-t", "--tools-bin-path", dest="tools_bin_path")
     parser.add_argument("-b", "--build-scripts-path", dest="build_scripts_path")
     parser.add_argument("-s", "--src-root", dest="src_root")
 
     options = parser.parse_args()

     utils = Utils()
     config = utils.jsonread(options.vmdk_config_path)

     print(options)

     disk_device = (utils.runshellcommand(
         "losetup --show -f {}".format(options.raw_image_path))).rstrip('\n')

     disk_partitions = utils.runshellcommand("kpartx -as {}".format(disk_device))
     device_name = disk_device.split('/')[2]
 
     if not os.path.exists(options.mount_path):
         os.mkdir(options.mount_path)

     loop_device_path = "/dev/mapper/{}p2".format(device_name)

 
     try:

         print("Generating PARTUUID for the loop device ...")
         partuuidval = (utils.runshellcommand(
             "blkid -s PARTUUID -o value {}".format(loop_device_path))).rstrip('\n')
         uuidval = (utils.runshellcommand(
             "blkid -s UUID -o value {}".format(loop_device_path))).rstrip('\n')
         if partuuidval == '':
             sgdiskout = utils.runshellcommand(
                 "sgdisk -i 2 {} ".format(disk_device))
             partuuidval = (re.findall(r'Partition unique GUID.*',
                                       sgdiskout))[0].split(':')[1].strip(' ').lower()
 
         if partuuidval == '':

             raise RuntimeError("Cannot generate partuuid")
 
         # Mount the loop device

         print("Mounting the loop device for customization ...")
         utils.runshellcommand(
             "mount -t ext4 {} {}".format(loop_device_path, options.mount_path))

         shutil.rmtree(options.mount_path + "/installer", ignore_errors=True)
         shutil.rmtree(options.mount_path + "/LOGS", ignore_errors=True)
         # Clear the root password if not set explicitly from the config file

         if config['password']['text'] != 'PASSWORD':
             utils.replaceinfile(options.mount_path + "/etc/shadow",
                                 'root:.*?:', 'root:*:')

         # Clear machine-id so it gets regenerated on boot
         open(options.mount_path + "/etc/machine-id", "w").close()
         os.remove(options.mount_path + "/etc/fstab")
 
         f = open(options.mount_path + "/etc/fstab", "w")

         if uuidval != '':

             f.write("UUID={}    /    ext4    defaults 1 1\n".format(uuidval))
         else:
             f.write("PARTUUID={}    /    ext4    defaults 1 1\n".format(partuuidval))

         f.close()

         utils.replaceinfile(options.mount_path + "/boot/grub/grub.cfg",
                             "rootpartition=PARTUUID=.*$",
                             "rootpartition=PARTUUID={}".format(partuuidval))

 
         if os.path.exists(options.additional_rpms_path):

             print("Installing additional rpms")

             os.mkdir(options.mount_path + "/additional_rpms")
             os.mkdir(options.mount_path + "/var/run")

             utils.copyallfiles(options.additional_rpms_path,
                                options.mount_path + "/additional_rpms")
             utils.runshellcommand(
                 "chroot {} /bin/bash -c 'rpm -i /additional_rpms/*'".format(options.mount_path))

             shutil.rmtree(options.mount_path + "/additional_rpms", ignore_errors=True)

             shutil.rmtree(options.additional_rpms_path, ignore_errors=True)

 
         utils.runshellcommand("mount -o bind /proc {}".format(options.mount_path + "/proc"))

         utils.runshellcommand("mount -o bind /dev {}".format(options.mount_path + "/dev"))
         utils.runshellcommand("mount -o bind /dev/pts {}".format(options.mount_path + "/dev/pts"))

         utils.runshellcommand("mount -o bind /sys {}".format(options.mount_path + "/sys"))
 
         if 'additionalfiles' in config:

             print("  Copying additional files into the raw image ...")

             for filetuples in config['additionalfiles']:

                 for src, dest in filetuples.items():
                     if (os.path.isdir(options.build_scripts_path + '/' +
                                       options.image_name + '/' + src)):
                         shutil.copytree(options.build_scripts_path + '/' +
                                         options.image_name + '/' + src,
                                         options.mount_path + dest, True)

                     else:

                         shutil.copyfile(options.build_scripts_path + '/' +
                                         options.image_name + '/' + src,
                                         options.mount_path + dest)

 
 
         if 'postinstallscripts' in config:

             print("  Running post install scripts ...")

             if not os.path.exists(options.mount_path + "/tempscripts"):
                 os.mkdir(options.mount_path + "/tempscripts")
             for script in config['postinstallscripts']:

                 shutil.copy(options.build_scripts_path + '/' +
                             options.image_name + '/' + script,
                             options.mount_path + "/tempscripts")

             for script in os.listdir(options.mount_path + "/tempscripts"):

                 print("     ...running script {}".format(script))
                 utils.runshellcommand(
                     "chroot {} /bin/bash -c '/tempscripts/{}'".format(options.mount_path, script))

             shutil.rmtree(options.mount_path + "/tempscripts", ignore_errors=True)
 

     except Exception as e:

         print(e)

 
     finally:

         utils.runshellcommand("umount -l {}".format(options.mount_path + "/sys"))

         utils.runshellcommand("umount -l {}".format(options.mount_path + "/dev/pts"))
         utils.runshellcommand("umount -l {}".format(options.mount_path + "/dev"))

         utils.runshellcommand("umount -l {}".format(options.mount_path + "/proc"))

 
         utils.runshellcommand("sync")

         utils.runshellcommand("umount -l {}".format(options.mount_path))
 

         mount_out = utils.runshellcommand("mount")

         print("List of mounted devices:")
         print(mount_out)

         utils.runshellcommand("kpartx -d {}".format(disk_device))
         utils.runshellcommand("losetup -d {}".format(disk_device))
 
         shutil.rmtree(options.mount_path)
 
         photon_release_ver = os.environ['PHOTON_RELEASE_VER']
         photon_build_num = os.environ['PHOTON_BUILD_NUM']
         raw_image = options.raw_image_path
         new_name = ""
         img_path = os.path.dirname(os.path.realpath(raw_image))
         # Rename gce image to disk.raw
         if options.image_name == "gce":

             print("Renaming the raw file to disk.raw ...")

             new_name = img_path + '/disk.raw'
 
         else:

             new_name = (img_path + '/photon-' + options.image_name +
                         '-' + photon_release_ver + '-' +
                         photon_build_num + '.raw')

 
         shutil.move(raw_image, new_name)
         raw_image = new_name
 
         if config['artifacttype'] == 'tgz':

             print("Generating the tar.gz artifact ...")
             tarname = (img_path + '/photon-' + options.image_name +
                        '-' + photon_release_ver + '-' +
                        photon_build_num + '.tar.gz')
 

             tgzout = tarfile.open(tarname, "w:gz")

             tgzout.add(raw_image, arcname=os.path.basename(raw_image))
             tgzout.close()

         elif config['artifacttype'] == 'xz':

             print("Generating the xz artifact ...")

             utils.runshellcommand("xz -z -k {}".format(raw_image))

 #            tarname = img_path + '/photon-' + options.image_name +
 #            '-' + photon_release_ver + '-' + photon_build_num +
 #            '.xz'

 #            tgzout = tarfile.open(tarname, "w:xz")
 #            tgzout.add(raw_image, arcname=os.path.basename(raw_image))
 #            tgzout.close()

         elif config['artifacttype'] == 'vhd':

             relrawpath = os.path.relpath(raw_image, options.src_root)

             vhdname = (os.path.dirname(relrawpath) + '/photon-' +
                        options.image_name + '-' + photon_release_ver + '-' +
                        photon_build_num + '.vhd')
             print("Converting raw disk to vhd ...")
             info_output = utils.runshellcommand(
                 "docker run -v {}:/mnt:rw anishs/qemu-img info -f raw --output json {}"
                 .format(options.src_root, '/mnt/' + relrawpath))

             mbsize = 1024 * 1024

             mbroundedsize = ((int(json.loads(info_output)["virtual-size"])/mbsize + 1) * mbsize)

             utils.runshellcommand(
                 "docker run -v {}:/mnt:rw anishs/qemu-img resize -f raw {} {}"
                 .format(options.src_root, '/mnt/' + relrawpath, mbroundedsize))
             utils.runshellcommand(
                 "docker run -v {}:/mnt:rw anishs/qemu-img convert {} -O "
                 "vpc -o subformat=fixed,force_size {}"
                 .format(options.src_root, '/mnt/' + relrawpath, '/mnt/' + vhdname))

         elif config['artifacttype'] == 'ova':

             create_ova_image(raw_image, options.tools_bin_path,
                              options.build_scripts_path + '/' + options.image_name, config)

             if 'customartifacts' in config:
                 if 'postinstallscripts' in config['customartifacts']:
                     custom_path = img_path + '/photon-custom'
                     if not os.path.exists(custom_path):
                         os.mkdir(custom_path)
                     index = 1
                     for script in config['customartifacts']['postinstallscripts']:

                         print("Creating custom ova {}...".format(index))

                         if index > 1:

                             raw_image_custom = (img_path + "/photon-custom-{}".format(index) +
                                                 photon_release_ver + '-' +
                                                 photon_build_num + '.raw')

                         else:

                             raw_image_custom = (img_path + "/photon-custom-" +
                                                 photon_release_ver + '-' +
                                                 photon_build_num + '.raw')

                         shutil.move(raw_image, raw_image_custom)

                         disk_device = (
                             utils.runshellcommand(
                                 "losetup --show -f {}".format(raw_image_custom))).rstrip('\n')
                         disk_partitions = utils.runshellcommand(
                             "kpartx -as {}".format(disk_device))

                         device_name = disk_device.split('/')[2]

                         loop_device_path = "/dev/mapper/{}p2".format(device_name)

                         print("Mounting the loop device for ova customization ...")
                         utils.runshellcommand(
                             "mount -t ext4 {} {}".format(loop_device_path, custom_path))

                         if not os.path.exists(custom_path + "/tempscripts"):
                             os.mkdir(custom_path + "/tempscripts")

                         shutil.copy(options.build_scripts_path + '/' + options.image_name +
                                     '/' + script, custom_path + "/tempscripts")
                         print("Running custom ova script {}".format(script))
                         utils.runshellcommand("chroot {} /bin/bash -c '/tempscripts/{}'"
                                               .format(custom_path, script))

                         shutil.rmtree(custom_path + "/tempscripts", ignore_errors=True)
                         utils.runshellcommand("umount -l {}".format(custom_path))

 
                         mount_out = utils.runshellcommand("mount")

                         print("List of mounted devices:")
                         print(mount_out)

                         utils.runshellcommand("kpartx -d {}".format(disk_device))
                         utils.runshellcommand("losetup -d {}".format(disk_device))

                         create_ova_image(raw_image_custom, options.tools_bin_path,
                                          options.build_scripts_path + '/' + options.image_name,
                                          config)

                         raw_image = raw_image_custom
                         index = index + 1
 

                     shutil.rmtree(custom_path)

 
         else:
             raise ValueError("Unknown output format")
 
         if config['keeprawdisk'] == 'false':

             os.remove(raw_image)