#!/usr/bin/python
#
# mic-chroot : chroot into an image or file system in order to change it
#
# Copyright 2009, Intel Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; version 2 of the License.
#
# 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 Library General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

import os
import sys
import optparse
import subprocess
import uuid
import re
import time
import mic.imgcreate as imgcreate

def error(msg):
    print("Error: %s" % msg)
    sys.exit(-1)

def launch_vm(imgfile):
    format = imgcreate.get_image_type(imgfile)
    if format in ("livecd", "liveusb", "raw"):
        qemu = "/usr/bin/kvm"
        if not os.path.exists(qemu):
            qemu = "/usr/bin/qemu-kvm"
            if not os.path.exists(qemu):
                qemu = "/usr/bin/qemu"
                if not os.path.exists(qemu):
                    error("Please install qemu or kvm beforehand.")

        bootdisk = {"livecd":"-cdrom", "liveusb":"-hda", "raw":"-hda"}
        bootsrc = {"livecd":"d", "liveusb":"c", "raw":"c"}
        args = [ qemu, bootdisk[format], imgfile, "-boot", bootsrc[format], "-m", "%d" % 512, "-std-vga", "-k", "en-us" ]
    elif format == "vmdk":
        vmplayer = "/usr/bin/vmplayer"
        if not os.path.exists(vmplayer):
            error("Please install vmplayer beforehand.")
        args = [ vmplayer, imgfile[0:-5] + ".vmx" ]
    elif format == "vdi":
        vboxmanage = "/usr/bin/VBoxManage"
        if not os.path.exists(vboxmanage):
            error("Please install VirtualBox beforehand.")

        """ Add root to group vboxusers """
        args = [ "/usr/sbin/usermod", "-G", "vboxusers", "-a", "root" ]
        p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        (output, errorout) = p.communicate()
        if p.returncode != 0:
            error("Failed to add root to group vboxusers")

        """ Create and register a VirtualBox VM """
        myuuid = "%s" % uuid.uuid4()
        args = [ vboxmanage, "createvm", "--register", "--name", myuuid, "--uuid", myuuid, "--ostype", "Linux" ]
        p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        (output, errorout) = p.communicate()
        if p.returncode != 0:
            error("Failed to create a VirtualBox VM: %s\n%s" % (output, errorout))

        """ Set some properties of this VM """
        args = [ vboxmanage, "modifyvm", myuuid, "--memory", "512", "--pae", "on", "--hda", imgfile ]
        p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        (output, errorout) = p.communicate()
        if p.returncode != 0:
            error("Failed to set properties of VirtualBox VM")

        """ Start this VirtualBox VM """
        args = [ vboxmanage, "startvm", myuuid ]
    else:
        error("Invalid image file.")

    try:
        if format in ("livecd", "liveusb", "raw"):
            subprocess.call(["/sbin/modprobe", "kvm_intel"])
        elif format in ("vmdk", "vdi"):
            subprocess.call(["/sbin/modprobe", "-r", "kvm_intel"])
        p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        (output, errorout) = p.communicate()
        print output
        print errorout
        if p.returncode != 0:
            if format in ("livecd", "liveusb", "raw"):
                # New kvm changed option -std-vga to "-vga std", so try new one
                args = [ qemu, bootdisk[format], imgfile, "-boot", bootsrc[format], "-m", "%d" % 512, "-vga", "std", "-k", "en-us" ]
                p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                (output, errorout) = p.communicate()
                if p.returncode != 0:
                    raise OSError(-p.returncode, errorout)
            else:    
                raise OSError(-p.returncode, errorout)
        if format == "vdi":
            poweroff = -3
            vmstate = "running"
            while vmstate != "powered off":
                args = [ vboxmanage, "showvminfo", myuuid ]
                p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                (output, errorout) = p.communicate()
                lines = output.split("\n")
                for line in lines:
                    if re.match(r"^State:.*powered off|State:.*saved", line):
                        """ Maybe it isn't started, need to check it in a few seconds """
                        poweroff = poweroff + 1
                        if poweroff >= 0:
                            vmstate = "powered off"
                        break
                time.sleep(2)

    except OSError, (err, msg):
        error("Failed to launch VM using %(imagefile)s: %(msg)s" % {"imagefile":imgfile, "msg":msg})
    finally:
        if format == "vdi":
            """ Discard VM state """
            args = [ vboxmanage, "discardstate", myuuid ]
            p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            (output, errorout) = p.communicate()

            """ Detach disk """
            args = [ vboxmanage, "modifyvm", myuuid, "--hda", "none" ]
            p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            (output, errorout) = p.communicate()

            """ Remove VM """
            args = [ vboxmanage, "unregistervm", myuuid, "--delete" ]
            p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            (output, errorout) = p.communicate()

            """ Close medium """
            args = [ vboxmanage, "closemedium", "disk", imgfile ]
            p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            (output, errorout) = p.communicate()


def parse_options(args):
    parser = optparse.OptionParser("Usage: %prog [options] [image file]")
    (options, args) = parser.parse_args()
    if len(args) != 1:
        error("Error: too few arguments.")
    if not os.path.isfile(args[0]):
        error("Error: invalid image file.")
    options.imgfile = os.path.abspath(os.path.expanduser(args[0]))

    return options

if __name__ == "__main__":
    if os.geteuid () != 0:
        error("You must run mic-vm-launcher as root")

    if not os.environ.has_key("DISPLAY") or not os.environ["DISPLAY"]:
        error("You must run mic-vm-launcher under X environment")

    if not os.environ.has_key("XAUTHORITY") or not os.environ["XAUTHORITY"]:
        error("You must run mic-vm-launcher under X environment")

    options = parse_options(sys.argv[1:])
    if options.imgfile:
        launch_vm(options.imgfile)
