#!/usr/bin/python2.6
#
# The Qubes OS Project, http://www.qubes-os.org
#
# Copyright (C) 2010  Joanna Rutkowska <joanna@invisiblethingslab.com>
#
# 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; either version 2
# of the License, or (at your option) any later version.
#
# 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#
#

from qubes.qubes import QubesVmCollection
from qubes.qubes import QubesException
from qubes.qubes import qubes_store_filename
from qubes.qubes import qubes_base_dir
from qubes.qubes import qubes_templates_dir
from qubes.qubes import qubes_appvms_dir
from optparse import OptionParser

import os
import time
import subprocess
import sys
import re

def size_to_human (size):
    if size < 1024:
        return str (size);
    elif size < 1024*1024:
        return str(round(size/1024.0,1)) + ' KiB'
    elif size < 1024*1024*1024:
        return str(round(size/(1024.0*1024),1)) + ' MiB'
    else:
        return str(round(size/(1024.0*1024*1024),1)) + ' GiB'

fields = {
    "qid": {"func": "vm.qid"},

    "name": {"func": "('=>' if backup_collection.get_default_template_vm() is not None\
             and vm.qid == backup_collection.get_default_template_vm().qid else '')\
             + ('[' if vm.is_template() else '')\
             + ('{' if vm.is_netvm() else '')\
             + vm.name \
             + (']' if vm.is_template() else '')\
             + ('}' if vm.is_netvm() else '')"},

    "type": {"func": "'Tpl' if vm.is_template() else \
             (' Net' if vm.is_netvm() else 'App')"},

    "updbl" : {"func": "'Yes' if vm.is_updateable() else ''"},

    "template": {"func": "'n/a' if vm.is_template() or vm.template_vm is None else\
                 find_template_name(backup_collection[vm.template_vm.qid].name,\
                         options.replace_template)"},

    "netvm": {"func": "'n/a' if vm.is_netvm() else\
              ('*' if vm.uses_default_netvm else '') +\
              backup_collection[vm.netvm_vm.qid].name\
                     if vm.netvm_vm is not None else '-'"},

    "label" : {"func" : "vm.label.name"},
}

def is_vm_included_in_backup (backup_dir, vm):
    if vm.qid == 0:
        # Dom0 is not included, obviously
        return False

    backup_vm_dir_path = vm.dir_path.replace (qubes_base_dir, backup_dir)

    if os.path.exists (backup_vm_dir_path):
        return True
    else:
        return False

def restore_vm_file (backup_dir, file_path):

    backup_file_path = file_path.replace (qubes_base_dir, backup_dir)
    #print "cp -rp {0} {1}".format (backup_file_path, file_path)

    # We prefer to use Linux's cp, because it nicely handles sparse files
    retcode = subprocess.call (["cp", "-p", backup_file_path, file_path])
    if retcode != 0:
        print "*** Error while copying file {0} to {1}".format(backup_file_path, file_path)
        exit (1)

def restore_vm_dir (backup_dir, src_dir, dst_dir):

    backup_src_dir = src_dir.replace (qubes_base_dir, backup_dir)

    # We prefer to use Linux's cp, because it nicely handles sparse files
    retcode = subprocess.call (["cp", "-rp", backup_src_dir, dst_dir])
    if retcode != 0:
        print "*** Error while copying file {0} to {1}".format(backup_src_dir, dest_dir)
        exit (1)

def find_template_name(template, replaces):
    rx_replace = re.compile("(.*):(.*)")
    for r in replaces:
        m = rx_replace.match(r)
        if m.group(1) == template:
            return m.group(2)

    return template

def main():
    usage = "usage: %prog [options] <backup-dir>"
    parser = OptionParser (usage)

    parser.add_option ("--skip-broken", action="store_true", dest="skip_broken", default=False,
                       help="Do not restore VMs that have missing templates or netvms")

    parser.add_option ("--ignore-missing", action="store_true", dest="ignore_missing", default=False,
                       help="Ignore missing templates or netvms, restore VMs anyway")

    parser.add_option ("--skip-conflicting", action="store_true", dest="skip_conflicting", default=False,
                       help="Do not restore VMs that are already present on the host")

    parser.add_option ("--force-root", action="store_true", dest="force_root", default=False,
                       help="Force to run, even with root privileges")

    parser.add_option ("--replace-template", action="append", dest="replace_template", default=[],
                       help="Restore VMs using another template, syntax: old-template-name:new-template-name (might be repeated)")

    (options, args) = parser.parse_args ()

    if (len (args) != 1):
        print "You must specify the backup directory (e.g. /mnt/backup/qubes-2010-12-01-235959)"
        exit (0)

    backup_dir = args[0]

    if not os.path.exists (backup_dir):
        print "The backup directory doesn't exist!"
        exit(1)

    backup_collection = QubesVmCollection(store_filename = backup_dir + "/qubes.xml")
    backup_collection.lock_db_for_reading()
    backup_collection.load()

    host_collection = QubesVmCollection()
    host_collection.lock_db_for_writing()
    host_collection.load()

    backup_vms_list = [vm for vm in backup_collection.values()]
    host_vms_list = [vm for vm in host_collection.values()]
    vms_to_restore = []

    fields_to_display = ["name", "type", "template", "updbl", "netvm", "label" ]

    # First calculate the maximum width of each field we want to display
    total_width = 0;
    for f in fields_to_display:
        fields[f]["max_width"] = len(f)
        for vm in backup_vms_list:
            l = len(str(eval(fields[f]["func"])))
            if l > fields[f]["max_width"]:
                fields[f]["max_width"] = l
        total_width += fields[f]["max_width"]

    print
    print "The following VMs are included in the backup:"
    print

    # Display the header
    s = ""
    for f in fields_to_display:
        fmt="{{0:-^{0}}}-+".format(fields[f]["max_width"] + 1)
        s += fmt.format('-')
    print s
    s = ""
    for f in fields_to_display:
        fmt="{{0:>{0}}} |".format(fields[f]["max_width"] + 1)
        s += fmt.format(f) 
    print s
    s = ""
    for f in fields_to_display:
        fmt="{{0:-^{0}}}-+".format(fields[f]["max_width"] + 1)
        s += fmt.format('-')
    print s

    there_are_conflicting_vms = False
    there_are_missing_templates = False
    there_are_missing_netvms = False
    # ... and the actual data
    for vm in backup_vms_list:
        if is_vm_included_in_backup (backup_dir, vm):
            s = ""
            good_to_go = True

            for f in fields_to_display:
                fmt="{{0:>{0}}} |".format(fields[f]["max_width"] + 1)
                s += fmt.format(eval(fields[f]["func"])) 

            if host_collection.get_vm_by_name (vm.name) is not None:
                s +=  " <-- A VM with the same name already exists on the host!"
                there_are_conflicting_vms = True
                good_to_go = False # Do not overwrite VMs on the host!

            if vm.template_vm is not None:
                templatevm_name = find_template_name(vm.template_vm.name, options.replace_template)
                template_vm_on_host = host_collection.get_vm_by_name (templatevm_name)

                # No template on the host?
                if not ((template_vm_on_host is not None) and template_vm_on_host.is_template()):

                    # Maybe the (custom) template is in the backup?
                    template_vm_on_backup = backup_collection.get_vm_by_name (templatevm_name)
                    if not ((template_vm_on_backup is not None) and template_vm_on_backup.is_template):
                        s += " <-- No matching template on the host or in the backup found!"
                        there_are_missing_templates = True
                        good_to_go = False if not (options.ignore_missing) else True
                
            if not vm.is_netvm() and vm.netvm_vm is not None:
                netvm_name = vm.netvm_vm.name
                netvm_vm_on_host = host_collection.get_vm_by_name (netvm_name)

                # No netv, on the host?
                if not ((netvm_vm_on_host is not None) and netvm_vm_on_host.is_netvm):

                    # Maybe the (custom) netvm is in the backup?
                    netvm_vm_on_backup = backup_collection.get_vm_by_name (netvm_name)
                    if not ((netvm_vm_on_backup is not None) and netvm_vm_on_backup.is_netvm):
                        s += " <-- No matching netvm on the host found!"
                        there_are_missing_netvms = True
                        good_to_go = False if not (options.ignore_missing) else True

            print s
            if good_to_go:
                vms_to_restore.append (vm)

    print

    if os.geteuid() == 0:
        print "*** Running this tool as root is strongly discouraged, this will lead you in permissions problems."
        if options.force_root:
            print "Continuing as commanded. You have been warned."
        else:
            print "Retry as unprivileged user."
            print "... or use --force-root to continue anyway."
            exit(1)

    if there_are_conflicting_vms:
        print "*** There VMs with conflicting names on the host! ***"
        if options.skip_conflicting:
            print "Those VMs will not be restored, the host VMs will not be overwritten!"
        else:
            print "Remove VMs with conflicting names from the host before proceeding."
            print "... or use --skip-conflicting to restore only those VMs that do not exist on the host."
            exit (1)

    print "The above VMs will be copied and added to your system."
    print "Exisiting VMs will not be removed."
 
    if there_are_missing_templates:
        print "*** One or more template VM is missing on the host! ***"
        if not (options.skip_broken or options.ignore_missing):
            print "Install it first, before proceeding with backup restore."
            print "Or pass: --skip-broken or --ignore-missing switch."
            exit (1)
        elif options.skip_broken:
            print "... VMs that depend on it will not be restored (--skip-broken used)"
        elif options.ignore_missing:
            print "... VMs that depend on it will be restored anyway (--ignore-missing used)"
        else:
            print "INTERNAL ERROR?!"
            exit (1)

    if there_are_missing_netvms:
        print "*** One or more network VM is missing on the host! ***"
        if not (options.skip_broken or options.ignore_missing):
            print "Install it first, before proceeding with backup restore."
            print "Or pass: --skip_broken or --ignore_missing switch."
            exit (1)
        elif options.skip_broken:
            print "... VMs that depend on it will not be restored (--skip-broken used)"
        elif options.ignore_missing:
            print "... VMs that depend on it be restored anyway (--ignore-missing used)"
        else:
            print "INTERNAL ERROR?!"
            exit (1)

    prompt = raw_input ("Do you want to proceed? [y/N] ")
    if not (prompt == "y" or prompt == "Y"):
        exit (0)

    # Add templates...
    for vm in [ vm for vm in vms_to_restore if vm.is_template()]:
        print "-> Restoring Template VM {0}...".format(vm.name)
        retcode = subprocess.call (["mkdir", "-p", vm.dir_path])
        if retcode != 0:
            print ("*** Cannot create directory: {0}?!".format(dest_dir))
            print ("Skiping...")
            continue

        restore_vm_dir (backup_dir, vm.dir_path, qubes_templates_dir);
        updateable = vm.updateable
        try:
            vm = host_collection.add_new_templatevm(vm.name, 
                                               conf_file=vm.conf_file,
                                               dir_path=vm.dir_path,
                                               installed_by_rpm=False)

            vm.updateable = updateable
            vm.verify_files()
        except Exception as err:
            print "ERROR: {0}".format(err)
            print "*** Skiping VM: {0}".vm.name
            if vm:
                host_collection.pop(vm.qid)
            continue

        try:
            vm.create_appmenus(verbose=True)
        except Exception as err:
            print "ERROR during appmenu restore: {0}".format(err)
            print "*** VM '{0}' will not have appmenus".format(vm.name)

    # ... then appvms...
    for vm in [ vm for vm in vms_to_restore if vm.is_appvm()]:

        print "-> Restoring AppVM {0}...".format(vm.name)
        retcode = subprocess.call (["mkdir", "-p", vm.dir_path])
        if retcode != 0:
            print ("*** Cannot create directory: {0}?!".format(dest_dir))
            print ("Skiping...")
            continue

        restore_vm_dir (backup_dir, vm.dir_path, qubes_appvms_dir);

        template_vm = None
        if vm.template_vm is not None:
            template_name = find_template_name(vm.template_vm.name, options.replace_template)
            template_vm = host_collection.get_vm_by_name(template_name)

        if not vm.uses_default_netvm:
            uses_default_netvm = False
            netvm_vm = host_collection.get_vm_by_name (vm.netvm_vm.name) if vm.netvm_vm is not None else None
        else:
            uses_default_netvm = True

        updateable = vm.updateable

        try:
            vm = host_collection.add_new_appvm(vm.name, template_vm,
                                          conf_file=vm.conf_file,
                                          dir_path=vm.dir_path,
                                          updateable=updateable,
                                          label=vm.label)
        except Exception as err:
            print "ERROR: {0}".format(err)
            print "*** Skiping VM: {0}".format(vm.name)
            if vm:
                host_collection.pop(vm.qid)
            continue

        if not uses_default_netvm:
            vm.uses_default_netvm = False
            vm.netvm_vm = netvm_vm

        try:
            vm.create_appmenus(verbose=True)
        except Exception as err:
            print "ERROR during appmenu restore: {0}".format(err)
            print "*** VM '{0}' will not have appmenus".format(vm.name)

        try:
            vm.verify_files()
        except Exception as err:
            print "ERROR: {0}".format(err)
            print "*** Skiping VM: {0}".format(vm.name)
            host_collection.pop(vm.qid)
            continue

    backup_collection.unlock_db()
    host_collection.save()
    host_collection.unlock_db()
    print "-> Done."
main()