#!/usr/bin/env python3
#
# Centos aarch64 image
#
# Copyright 2020 Linaro
#
# Authors:
#  Robert Foley <robert.foley@linaro.org>
#  Originally based on ubuntu.aarch64
#
# This code is licensed under the GPL version 2 or later.  See
# the COPYING file in the top-level directory.
#

import os
import sys
import subprocess
import basevm
import time
import traceback
import aarch64vm

DEFAULT_CONFIG = {
    'cpu'          : "max",
    'machine'      : "virt,gic-version=max",
    'install_cmds' : "yum install -y make ninja-build git python3 gcc gcc-c++ flex bison, "\
        "yum install -y glib2-devel pixman-devel zlib-devel, "\
        "yum install -y perl-Test-Harness, "\
        "alternatives --set python /usr/bin/python3, "\
        "sudo dnf config-manager "\
        "--add-repo=https://download.docker.com/linux/centos/docker-ce.repo,"\
        "sudo dnf install --nobest -y docker-ce.aarch64,"\
        "systemctl enable docker",
    # We increase beyond the default time since during boot
    # it can take some time (many seconds) to log into the VM.
    'ssh_timeout'  : 60,
}

class CentosAarch64VM(basevm.BaseVM):
    name = "centos.aarch64"
    arch = "aarch64"
    login_prompt = "localhost login:"
    prompt = '[root@localhost ~]#'
    image_name = "CentOS-8-aarch64-1905-dvd1.iso"
    image_link = "http://mirrors.usc.edu/pub/linux/distributions/centos/8.0.1905/isos/aarch64/"
    image_link += image_name
    BUILD_SCRIPT = """
        set -e;
        cd $(mktemp -d);
        sudo chmod a+r /dev/vdb;
        tar --checkpoint=.10 -xf /dev/vdb;
        ./configure {configure_opts};
        make --output-sync {target} -j{jobs} {verbose};
    """
    def set_key_perm(self):
        """Set permissions properly on certain files to allow
           ssh access."""
        self.console_wait_send(self.prompt,
                               "/usr/sbin/restorecon -R -v /root/.ssh\n")
        self.console_wait_send(self.prompt,
                "/usr/sbin/restorecon -R -v "\
                "/home/{}/.ssh\n".format(self._config["guest_user"]))

    def create_kickstart(self):
        """Generate the kickstart file used to generate the centos image."""
        # Start with the template for the kickstart.
        ks_file = self._source_path + "/tests/vm/centos-8-aarch64.ks"
        subprocess.check_call("cp {} ./ks.cfg".format(ks_file), shell=True)
        # Append the ssh keys to the kickstart file
        # as the post processing phase of installation.
        with open("ks.cfg", "a") as f:
            # Add in the root pw and guest user.
            rootpw = "rootpw --plaintext {}\n"
            f.write(rootpw.format(self._config["root_pass"]))
            add_user = "user --groups=wheel --name={} "\
                       "--password={} --plaintext\n"
            f.write(add_user.format(self._config["guest_user"],
                                    self._config["guest_pass"]))
            # Add the ssh keys.
            f.write("%post --log=/root/ks-post.log\n")
            f.write("mkdir -p /root/.ssh\n")
            addkey = 'echo "{}" >> /root/.ssh/authorized_keys\n'
            addkey_cmd = addkey.format(self._config["ssh_pub_key"])
            f.write(addkey_cmd)
            f.write('mkdir -p /home/{}/.ssh\n'.format(self._config["guest_user"]))
            addkey = 'echo "{}" >> /home/{}/.ssh/authorized_keys\n'
            addkey_cmd = addkey.format(self._config["ssh_pub_key"],
                                       self._config["guest_user"])
            f.write(addkey_cmd)
            f.write("%end\n")
        # Take our kickstart file and create an .iso from it.
        # The .iso will be provided to qemu as we boot
        # from the install dvd.
        # Anaconda will recognize the label "OEMDRV" and will
        # start the automated installation.
        gen_iso_img = 'genisoimage -output ks.iso -volid "OEMDRV" ks.cfg'
        subprocess.check_call(gen_iso_img, shell=True)

    def wait_for_shutdown(self):
        """We wait for qemu to shutdown the VM and exit.
           While this happens we display the console view
           for easier debugging."""
        # The image creation is essentially done,
        # so whether or not the wait is successful we want to
        # wait for qemu to exit (the self.wait()) before we return.
        try:
            self.console_wait("reboot: Power down")
        except Exception as e:
            sys.stderr.write("Exception hit\n")
            if isinstance(e, SystemExit) and e.code == 0:
                return 0
            traceback.print_exc()
        finally:
            self.wait()

    def build_base_image(self, dest_img):
        """Run through the centos installer to create
           a base image with name dest_img."""
        # We create the temp image, and only rename
        # to destination when we are done.
        img = dest_img + ".tmp"
        # Create an empty image.
        # We will provide this as the install destination.
        qemu_img_create = "qemu-img create {} 50G".format(img)
        subprocess.check_call(qemu_img_create, shell=True)

        # Create our kickstart file to be fed to the installer.
        self.create_kickstart()
        # Boot the install dvd with the params as our ks.iso
        os_img = self._download_with_cache(self.image_link)
        dvd_iso = "centos-8-dvd.iso"
        subprocess.check_call(["cp", "-f", os_img, dvd_iso])
        extra_args = "-cdrom ks.iso"
        extra_args += " -drive file={},if=none,id=drive1,cache=writeback"
        extra_args += " -device virtio-blk,drive=drive1,bootindex=1"
        extra_args = extra_args.format(dvd_iso).split(" ")
        self.boot(img, extra_args=extra_args)
        self.console_wait_send("change the selection", "\n")
        # We seem to need to hit esc (chr(27)) twice to abort the
        # media check, which takes a long time.
        # Waiting a bit seems to be more reliable before hitting esc.
        self.console_wait("Checking")
        time.sleep(5)
        self.console_wait_send("Checking", chr(27))
        time.sleep(5)
        self.console_wait_send("Checking", chr(27))
        print("Found Checking")
        # Give sufficient time for the installer to create the image.
        self.console_init(timeout=7200)
        self.wait_for_shutdown()
        os.rename(img, dest_img)
        print("Done with base image build: {}".format(dest_img))

    def check_create_base_img(self, img_base, img_dest):
        """Create a base image using the installer.
           We will use the base image if it exists.
           This helps cut down on install time in case we
           need to restart image creation,
           since the base image creation can take a long time."""
        if not os.path.exists(img_base):
            print("Generate new base image: {}".format(img_base))
            self.build_base_image(img_base);
        else:
            print("Use existing base image: {}".format(img_base))
        # Save a copy of the base image and copy it to dest.
        # which we will use going forward.
        subprocess.check_call(["cp", img_base, img_dest])

    def boot(self, img, extra_args=None):
        aarch64vm.create_flash_images(self._tmpdir, self._efi_aarch64)
        default_args = aarch64vm.get_pflash_args(self._tmpdir)
        if extra_args:
            extra_args.extend(default_args)
        else:
            extra_args = default_args
        # We always add these performance tweaks
        # because without them, we boot so slowly that we
        # can time out finding the boot efi device.
        if '-smp' not in extra_args and \
           '-smp' not in self._config['extra_args'] and \
           '-smp' not in self._args:
            # Only add if not already there to give caller option to change it.
            extra_args.extend(["-smp", "8"])
        # We have overridden boot() since aarch64 has additional parameters.
        # Call down to the base class method.
        super(CentosAarch64VM, self).boot(img, extra_args=extra_args)

    def build_image(self, img):
        img_tmp = img + ".tmp"
        self.check_create_base_img(img + ".base", img_tmp)

        # Boot the new image for the first time to finish installation.
        self.boot(img_tmp)
        self.console_init()
        self.console_wait_send(self.login_prompt, "root\n")
        self.console_wait_send("Password:",
                               "{}\n".format(self._config["root_pass"]))

        self.set_key_perm()
        self.console_wait_send(self.prompt, "rpm -q centos-release\n")
        enable_adapter = "sed -i 's/ONBOOT=no/ONBOOT=yes/g'" \
                         " /etc/sysconfig/network-scripts/ifcfg-enp0s1\n"
        self.console_wait_send(self.prompt, enable_adapter)
        self.console_wait_send(self.prompt, "ifup enp0s1\n")
        self.console_wait_send(self.prompt,
                               'echo "qemu  ALL=(ALL) NOPASSWD:ALL" | '\
                               'sudo tee /etc/sudoers.d/qemu\n')
        self.console_wait(self.prompt)

        # Rest of the commands we issue through ssh.
        self.wait_ssh(wait_root=True)

        # If the user chooses *not* to do the second phase,
        # then we will jump right to the graceful shutdown
        if self._config['install_cmds'] != "":
            install_cmds = self._config['install_cmds'].split(',')
            for cmd in install_cmds:
                self.ssh_root(cmd)
        self.ssh_root("poweroff")
        self.wait_for_shutdown()
        os.rename(img_tmp, img)
        print("image creation complete: {}".format(img))
        return 0

if __name__ == "__main__":
    defaults = aarch64vm.get_config_defaults(CentosAarch64VM, DEFAULT_CONFIG)
    sys.exit(basevm.main(CentosAarch64VM, defaults))