#!/usr/bin/env python
# This file is part of Cockpit.
#
# Copyright (C) 2015 Red Hat, Inc.
#
# Cockpit is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 2.1 of the License, or
# (at your option) any later version.
#
# Cockpit is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Cockpit; If not, see <http://www.gnu.org/licenses/>.

# vm-create  -- Make a root image suitable for use with vm-run.
#
# Usage: vm-create
#
# Installs the OS indicated by TEST_OS into the image
# for test machine and tweaks it to be useable with
# vm-run and testlib.py.

import argparse
import os
import shutil
import subprocess
import sys
import time

from common import testvm
from common import testinfra

BASE = os.path.dirname(__file__)

parser = argparse.ArgumentParser(description='Create a virtual machine image')
parser.add_argument('-v', '--verbose', action='store_true', help='Display verbose progress details')
parser.add_argument('-s', '--sit', action='store_true', help='Sit and wait if setup script fails')
parser.add_argument('-n', '--no-save', action='store_true', help='Don\'t save the new image')
parser.add_argument('-u', '--upload', action='store_true', help='Upload the image after creation')
parser.add_argument('--no-build', action='store_true', dest='no_build',
                    help='Don''t build packages and create the vm without build capabilities')
parser.add_argument("--store", default=None, help="Where to send images")
parser.add_argument('image', nargs='?', default=testinfra.DEFAULT_IMAGE, help='The image to create')
args = parser.parse_args()

class MachineBuilder:
    def __init__(self, machine):
        self.machine = machine
        self.macaddr = None

    def bootstrap_system(self):
        assert not self.machine._domain

        if not os.path.exists(self.machine.run_dir):
            os.makedirs(self.machine.run_dir, 0750)

        bootstrap_script = os.path.join(testinfra.TEST_DIR, "images", "scripts", "%s.bootstrap" % (self.machine.image, ))

        if os.path.exists(self.machine.image_file):
            os.unlink(self.machine.image_file)

        if os.path.isfile(bootstrap_script):
            subprocess.check_call([ bootstrap_script, self.machine.image_file ])
        else:
            raise testvm.Failure("Unsupported OS %s: %s not found." % (self.machine.image, bootstrap_script))

    def run_setup_script(self, script, source):
        """Prepare a test image further by running some commands in it."""
        self.machine.start(maintain=True, macaddr=self.macaddr)
        try:
            self.macaddr = self.machine.macaddr
            self.machine.wait_boot(wait_for_running_timeout=120)
            self.machine.upload([ os.path.join(testinfra.TEST_DIR, "images", "scripts", "lib") ], "/var/lib/testvm")
            self.machine.upload([script], "/var/tmp/SETUP")
            if source:
                self.machine.upload([source], "/var/tmp")
            self.machine.upload([ os.path.join(testinfra.TEST_DIR, "images", "scripts", "lib", "base") ],
                                "/var/tmp/cockpit-base")

            if "rhel" in self.machine.image:
                self.machine.upload([ os.path.expanduser("~/.rhel") ], "/root/")

            env = {
                "TEST_OS": self.machine.image
            }
            if source:
                env["TEST_SOURCE"] = os.path.join("/var/tmp", os.path.basename(source))
            self.machine.message("run setup script on guest")

            try:
                self.machine.execute(script="/var/tmp/SETUP", environment=env, quiet=not self.machine.verbose)
                self.machine.execute(command="rm -f /var/tmp/SETUP")
                self.machine.execute(command="rm -rf /root/.rhel")

                if self.machine.image == 'openshift':
                    # update our local openshift kube config file to match the new image
                    self.machine.download("/root/.kube/config", "verify/files/openshift.kubeconfig")

            except subprocess.CalledProcessError, ex:
                if args.sit:
                    sys.stderr.write("ADDRESS: {0}\n".format(self.machine.address))
                    raw_input ("Press RET to continue... ")
                raise testvm.Failure("setup failed with code {0}\n".format(ex.returncode))

        finally:
            self.machine.stop(timeout_sec=60)

    def boot_system(self):
        """Start the system to make sure it can boot, then shutdown cleanly
        This also takes care of any selinux relabeling setup triggered
        Don't wait for an ip address during start, since the system might reboot"""
        self.machine.start(maintain=True, macaddr=self.macaddr, wait_for_ip=False)
        self.machine.reset_reboot_flag()
        try:
            self.macaddr = self.machine.macaddr
            # it could be that selinux relabeling needs a reboot
            self.machine.wait_boot(wait_for_running_timeout=120, allow_one_reboot=True)
        finally:
            self.machine.stop(timeout_sec=120)

    def build(self):
        self.bootstrap_system()

        # gather the scripts, separated by reboots
        script = os.path.join(testinfra.TEST_DIR, "images", "scripts", "%s.setup" % (self.machine.image, ))

        if not os.path.exists(script):
            raise testvm.Failure("Unsupported image %s: %s not found." % (self.machine.image, script))

        self.machine.message("Running setup script %s" % (script))
        if args.no_build:
            source = None
        else:
            source = subprocess.check_output([ "../tools/make-source" ]).strip()
        self.run_setup_script(script, source)

        tries_left = 3
        successfully_booted = False
        while tries_left > 0:
            try:
                # make sure we can boot the system
                self.boot_system()
                successfully_booted = True
                break
            except:
                # we might need to wait for the image to become available again
                # accessing it in maintain=True mode successively can trigger qemu errors
                time.sleep(3)
                tries_left -= 1
        if not successfully_booted:
            raise testvm.Failure("Unable to verify that machine boot works.")

    def save(self):
        data_dir = os.path.join(os.environ.get("TEST_DATA", testinfra.TEST_DIR), "images")
        images_dir = os.path.join(testinfra.TEST_DIR, "images")

        if not os.path.exists(data_dir):
            os.makedirs(data_dir, 0750)

        if not os.path.exists(self.machine.image_file):
            raise testvm.Failure("Nothing to save.")

        partial = os.path.join(data_dir, self.machine.image + ".partial")

        # Copy image via convert, to make it sparse again
        subprocess.check_call([ "qemu-img", "convert", "-c", "-O", "qcow2", self.machine.image_file, partial ])

        # Hash the image here
        (sha, x1, x2) = subprocess.check_output([ "sha1sum", partial ]).partition(" ")
        if not sha:
            raise testvm.Failure("sha1sum returned invalid output")

        name = self.machine.image + "-" + sha + ".qcow2"
        data_file = os.path.join(data_dir, name)
        shutil.move(partial, data_file)

        # Update the images symlink
        if os.path.islink(self.machine.image_base):
            os.unlink(self.machine.image_base)
        os.symlink(name, self.machine.image_base)

        # Handle alternate TEST_DATA
        image_file = os.path.join(images_dir, name)
        if not os.path.exists(image_file):
            os.symlink(os.path.abspath(data_file), image_file)

try:
    os.chdir(BASE)
    testvm.VirtMachine.memory_mb = 2048
    machine = testvm.VirtMachine(verbose=args.verbose, image=args.image, fetch=False)
    builder = MachineBuilder(machine)
    builder.build()
    if not args.no_save:
        print "Saving..."
        builder.save()
        if args.upload:
            print "Uploading..."
            cmd = [ "./vm-upload" ]
            if args.store:
                cmd += [ "--store", args.store ]
            cmd += [ args.image ]
            subprocess.check_call(cmd)

except testvm.Failure, ex:
    print >> sys.stderr, "vm-create:", ex
    sys.exit(1)
