#!/usr/bin/env python3

"""
Script to test virtualization functionality

Copyright (C) 2013 Canonical Ltd.

Authors
  Jeff Marcom <jeff.marcom@canonical.com>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License version 3,
as published by the Free Software Foundation.

This program 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 General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""

from argparse import ArgumentParser
import configparser
import os
import logging
import lsb_release
import signal
from subprocess import Popen, PIPE
import sys
import tempfile
import time
import urllib.request


class XENTest(object):
    pass


class KVMTest(object):

    def __init__(self, image=None, timeout=500, debug_file="virt_debug"):
        self.image = image
        self.timeout = timeout
        self.debug_file = debug_file

    @classmethod
    def download_image(cls):
        """
        Downloads Cloud image for same release as host machine
        """

        # Check Ubuntu release info. Example {quantal, precise}
        release = lsb_release.get_lsb_information()["CODENAME"]

        # Construct URL
        cloud_url = "http://cloud-images.ubuntu.com"
        cloud_iso = release + "-server-cloudimg-i386-disk1.img"
        image_url = "/".join((
            cloud_url, release, "current", cloud_iso))

        logging.debug("Downloading {}, from {}".format(cloud_iso, cloud_url))

        # Attempt download
        try:
            resp = urllib.request.urlretrieve(image_url, cloud_iso)
        except urllib.error.HTTPError as exception:
            logging.error("Failed download {}: {}".format(image_url, exception))
            return False

        if not os.path.isfile(cloud_iso):
            return False

        return cloud_iso

    def boot_image(self, data_disk):
        """
        Attempts to boot the newly created qcow image using
        the config data defined in config.iso. Return instance
        of currently running virtual machine
        """

        logging.debug("Attempting boot for:{}".format(data_disk))

        # Set Arbitrary IP values
        netrange = "10.0.0.0/8"
        image_ip = "10.0.0.1"
        hostfwd = "tcp::2222-:22"

        params = \
        '''
        kvm -m {0} -net nic -net user,net={1},host={2},
        hostfwd={3} -drive file={4},if=virtio -display none -nographic
        '''.format(
            "256",
            netrange,
            image_ip,
            hostfwd,
            data_disk).replace("\n", "").replace("  ", "")

        logging.debug("Using params:{}".format(params))

        # Default file location for log file is in checkbox output directory
        checkbox_dir = os.getenv("CHECKBOX_DATA")

        if checkbox_dir != None:
            self.debug_file = os.path.join(checkbox_dir, self.debug_file)

        # Open VM STDERR/STDOUT log file for writing
        try:
            file = open(self.debug_file, 'w')
        except IOError:
            logging.error("Failed creating file:{}".format(self.debug_file))
            return False

        # Start Virtual machine
        self.process = Popen(params, stderr=file,
            stdout=file, universal_newlines=True, shell=True)

    def start(self):
        status = 1
        # Create temp directory:
        with tempfile.TemporaryDirectory("_kvm_test", 
        time.strftime("%b_%d_%Y_")) as temp_dir:
            
            os.chdir(temp_dir)
            if self.image is None:
                # Download cloud image
                self.image = self.download_image()

            if os.path.isfile(self.image):
 
                # Boot Virtual Machine
                instance = self.boot_image(self.image)

                if instance is not False:
                    time.sleep(self.timeout)
                    # Check to be sure VM boot was successful
                    if "END SSH HOST KEY KEYS" \
                    in open(self.debug_file, 'r').read():
                        print("Booted successfully", file=sys.stderr)
                        status = 0
                    else:
                        print("KVM instance failed to boot", file=sys.stderr)
                    self.process.terminate()
            else:
                print("Could not find: {}".format(self.image), file=sys.stderr)

        return status  
     

def test_kvm(args):
    print("Executing KVM Test", file=sys.stderr)

    DEFAULT_CFG = "/etc/checkbox.d/virtualization.cfg"
    image = ""
    timeout = ""

    config_file = DEFAULT_CFG
    config = configparser.SafeConfigParser()

    try:
        config.readfp(open(config_file))
        timeout = config.get("KVM", "timeout")
        image = config.get("KVM", "image") 
    except IOError:
        logging.warn("No config file found")
    except Exception as exception:
        logging.warn(exception)

    if image == "":
        image = args.image
    if timeout == "":
        timeout = args.timeout

    kvm_test = KVMTest(image, timeout)
    result = kvm_test.start()

    sys.exit(result)


def main():

    parser = ArgumentParser(description="Virtualization Test")
    subparsers = parser.add_subparsers()

    # Main cli options
    kvm_test_parser = subparsers.add_parser('kvm',
        help=("Run kvm virtualization test"))

    #xen_test_parser = subparsers.add_parser('xen',
    #    help=("Run xen virtualization test"))

    # Sub test options
    kvm_test_parser.add_argument('-i', '--image',
        type=str, default=None)
    kvm_test_parser.add_argument('-t', '--timeout',
        type=int, default=500)
    kvm_test_parser.add_argument('--debug',
        action="store_true")
    kvm_test_parser.set_defaults(func=test_kvm)

    args = parser.parse_args()

    if args.debug:
        logging.basicConfig(level=logging.DEBUG)

    args.func(args)

    
    

if __name__ == "__main__":
    main();
