#!/usr/bin/python2
# -*- encoding: utf8 -*-
#
# 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 multiprocessing import Event

from qubes.qubes import QubesVmCollection
from qubes.qubes import QubesException
from qubes.backup import backup_restore_header
from qubes.backup import backup_restore_prepare
from qubes.backup import backup_restore_print_summary
from qubes.backup import backup_restore_do
import qubes.backup
import sys
from optparse import OptionParser
from locale import getpreferredencoding

import os
import sys
import getpass

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

    parser.add_option ("--verify-only", action="store_true",
                       dest="verify_only", default=False,
                       help="Do not restore the data, only verify backup "
                            "integrify.")

    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 ("--rename-conflicting", action="store_true",
                       dest="rename_conflicting", default=False,
                       help="Restore VMs that are already present on the host under different name")

    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 (may be "
                            "repeated)")

    parser.add_option ("-x", "--exclude", action="append", dest="exclude", default=[],
                       help="Skip restore of specified VM (may be repeated)")

    parser.add_option ("--skip-dom0-home", action="store_false", dest="dom0_home", default=True,
                       help="Do not restore dom0 user home dir")

    parser.add_option ("--ignore-username-mismatch", action="store_true", dest="ignore_username_mismatch", default=False,
                       help="Ignore dom0 username mismatch while restoring homedir")

    parser.add_option ("-d", "--dest-vm", action="store", dest="appvm",
                       help="The AppVM to send backups to")

    parser.add_option ("-e", "--encrypted", action="store_true", dest="decrypt", default=False,
                       help="The backup is encrypted")

    parser.add_option ("-p", "--passphrase-file", action="store",
                       dest="pass_file", default=None,
                       help="File containing the pass phrase to use, or '-' to read it from stdin")

    parser.add_option ("-z", "--compressed", action="store_true", dest="compressed", default=False,
                       help="The backup is compressed")

    parser.add_option ("--debug", action="store_true", dest="debug",
                       default=False, help="Enable (a lot of) debug output")

    (options, args) = parser.parse_args ()

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

    backup_dir = args[0]
    vmlist = args[1:]

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

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

    restore_options = {}
    if options.ignore_missing:
        restore_options['use-default-template'] = True
        restore_options['use-default-netvm'] = True
    if options.replace_template:
        restore_options['replace-template'] = options.replace_template
    if options.rename_conflicting:
        restore_options['rename-conflicting'] = True
    if not options.dom0_home:
        restore_options['dom0-home'] = False
    if options.ignore_username_mismatch:
        restore_options['ignore-username-mismatch'] = True
    if options.exclude:
        restore_options['exclude'] = options.exclude
    if options.verify_only:
        restore_options['verify-only'] = True
    if options.debug:
        qubes.backup.BACKUP_DEBUG = True

    appvm = None
    if options.appvm is not None:
        appvm = host_collection.get_vm_by_name(options.appvm)
        if appvm is None:
            print >>sys.stderr, "ERROR: VM {0} does not exist".format(options.appvm)
            exit(1)

    if options.pass_file is not None:
        f = open(options.pass_file) if options.pass_file != "-" else sys.stdin
        passphrase = f.readline().rstrip()
        if f is not sys.stdin:
            f.close()
    else:
        passphrase = getpass.getpass("Please enter the pass phrase to decrypt/verify the backup: ")

    encoding = sys.stdin.encoding or getpreferredencoding()
    passphrase = passphrase.decode(encoding)

    print >> sys.stderr, "Checking backup content..."

    error_detected = Event()
    def error_callback(message):
        error_detected.set()
        print >> sys.stderr, message

    restore_info = None
    try:
        restore_info = backup_restore_prepare(
                backup_dir,
                passphrase=passphrase,
                options=restore_options,
                host_collection=host_collection,
                encrypted=options.decrypt,
                compressed=options.compressed,
                appvm=appvm,
                error_callback=error_callback)
    except QubesException as e:
        print >> sys.stderr, "ERROR: %s" % str(e)
        exit(1)

    if len(vmlist) > 0:
        for vm in restore_info.keys():
            if vm.startswith('$'):
                continue
            if not vm in vmlist:
                restore_info.pop(vm)

    backup_restore_print_summary(restore_info)

    there_are_conflicting_vms = False
    there_are_missing_templates = False
    there_are_missing_netvms = False
    dom0_username_mismatch = False

    for vm_info in restore_info.values():
        if 'excluded' in vm_info and vm_info['excluded']:
            continue
        if 'missing-template' in vm_info.keys():
            there_are_missing_templates = True
        if 'missing-netvm' in vm_info.keys():
            there_are_missing_netvms = True
        if 'already-exists' in vm_info.keys():
            there_are_conflicting_vms = True
        if 'username-mismatch' in vm_info.keys():
            dom0_username_mismatch = True

    print

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

    if there_are_conflicting_vms:
        print >> sys.stderr, "*** There VMs with conflicting names on the host! ***"
        if options.skip_conflicting:
            print >> sys.stderr, "Those VMs will not be restored, the host VMs will not be overwritten!"
        else:
            print >> sys.stderr, "Remove VMs with conflicting names from the host before proceeding."
            print >> sys.stderr, "... or use --skip-conflicting to restore only those VMs that do not exist on the host."
            print >> sys.stderr, "... or use --rename-conflicting to " \
                                 "restore those VMs under modified " \
                                 "name (with number at the end)"
            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 >> sys.stderr, "*** One or more template VM is missing on the host! ***"
        if not (options.skip_broken or options.ignore_missing):
            print >> sys.stderr, "Install it first, before proceeding with backup restore."
            print >> sys.stderr, "Or pass: --skip-broken or --ignore-missing switch."
            exit (1)
        elif options.skip_broken:
            print >> sys.stderr, "... VMs that depend on it will not be restored (--skip-broken used)"
        elif options.ignore_missing:
            print >> sys.stderr, "... VMs that depend on it will be restored anyway (--ignore-missing used)"
        else:
            print >> sys.stderr, "INTERNAL ERROR?!"
            exit (1)

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

    if 'dom0' in restore_info.keys() and options.dom0_home:
        if dom0_username_mismatch:
            print >> sys.stderr, "*** Dom0 username mismatch! This can break some settings ***"
            if not options.ignore_username_mismatch:
                print >> sys.stderr, "Skip dom0 home restore (--skip-dom0-home)"
                print >> sys.stderr, "Or pass: --ignore-username-mismatch to continue anyway"
                exit(1)
            else:
                print >> sys.stderr, "Continuing as directed"
        print >> sys.stderr, "While restoring user homedir, existing files/dirs will be backed up in 'home-pre-restore-<current-time>' dir"

    if options.pass_file is None:
        if raw_input("Do you want to proceed? [y/N] ").upper() != "Y":
            exit(0)

    try:
        backup_restore_do(restore_info,
                host_collection=host_collection,
                error_callback=error_callback)
    except QubesException as e:
        print >> sys.stderr, "ERROR: %s" % str(e)

    host_collection.unlock_db()

    if error_detected.is_set():
        print "-> Completed with errors!"
        exit(1)
    else:
        print "-> Done."
main()