#!/usr/bin/python2
 
 import os
 import re
 import shutil
 import tarfile
 import fileinput
 from optparse import OptionParser
 from utils import Utils

 import json

 
 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() == '</VirtualSystem>':
                     for ovfinfoline in lines:
                         print ovfinfoline,
                 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 = OptionParser(usage)
 
     parser.add_option("-r", "--raw-image-path",  dest="raw_image_path")
     parser.add_option("-c", "--vmdk-config-path", dest="vmdk_config_path")
     parser.add_option("-w",  "--working-directory",  dest="working_directory")
     parser.add_option("-m",  "--mount-path",  dest="mount_path")
     parser.add_option("-a",  "--additional-rpms-path",  dest="additional_rpms_path")
     parser.add_option("-i",  "--image-name",  dest="image_name")
     parser.add_option("-t",  "--tools-bin-path",  dest="tools_bin_path")
     parser.add_option("-b",  "--build-scripts-path",  dest="build_scripts_path")

     parser.add_option("-s",  "--src-root",  dest="src_root")

 
     (options,  args) = 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(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(additional_rpms_path, ignore_errors=True)
 
         utils.runshellcommand("mount -o bind /proc {}".format(options.mount_path + "/proc"))

         utils.runshellcommand("mount -o rbind /dev {}".format(options.mount_path + "/dev"))
         utils.runshellcommand("mount --make-rslave {}".format(options.mount_path + "/dev"))

         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.iteritems():
                     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)
 
         utils.runshellcommand("umount -l {}".format(options.mount_path + "/sys"))

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

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

         mount_out = utils.runshellcommand("mount")
         print "List of mounted devices:"
         print mount_out
 

     finally:
         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'] == '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)