#!/usr/bin/python2
#
# The Qubes OS Project, http://www.qubes-os.org
#
# Copyright (C) 2011  Marek Marczykowski <marmarek@mimuw.edu.pl>
#
# 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 optparse import OptionParser
import subprocess
import os
import time
import glob
import sys

def main():
    usage = "usage: %prog [options] <template-name>"
    parser = OptionParser (usage)
    parser.add_option ("--force", action="store_true", dest="force", default=False,
            help="Do not prompt for comfirmation")

    (options, args) = parser.parse_args ()
    if (len (args) != 1):
        parser.error ("You must specify TemplateVM name!")
    vmname = args[0]

    if os.geteuid() != 0:
        print >> sys.stderr, "ERROR: This tool must be run as root!"
        exit(1)

    qvm_collection = QubesVmCollection()
    qvm_collection.lock_db_for_reading()
    qvm_collection.load()
    qvm_collection.unlock_db()

    vm = qvm_collection.get_vm_by_name(vmname)
    if vm is None:
        print >> sys.stderr, "A VM with the name '{0}' does not exist in the system.".format(vmname)
        exit(1)

    if not vm.is_template():
        print >> sys.stderr, "A VM '{0}' is not template.".format(vmname)
        exit(1)

    if vm.is_running():
        print >> sys.stderr, "You must stop VM first."
        exit(1)

    oldcow_img = vm.rootcow_img + '.old'
    oldcow_stat = os.stat(oldcow_img)
    oldcow_time_str = time.strftime("%F %T", time.gmtime(oldcow_stat.st_mtime))

    root_stat = os.stat(vm.root_img)
    old_dmdev = "/dev/mapper/snapshot-{0:x}:{1}-{2:x}:{3}".format(
        root_stat[2], root_stat[1],
        oldcow_stat[2], oldcow_stat[1])

    snapshots = glob.glob('/dev/mapper/snapshot-{0:x}:{1}-*'.format(root_stat[2], root_stat[1]))
    snapshot_present = False
    for dev in snapshots:
        if dev == old_dmdev:
            snapshot_present = True
        else:
            print >> sys.stderr, "ERROR: You must shutdown all VMs running system older/newer than last good one."
            exit(1)

    root_blocks = os.path.getsize(vm.root_img)/512
    if not snapshot_present:
        p = subprocess.Popen (["/etc/xen/scripts/block-snapshot", "prepare",
            "snapshot", "{0}:{1}".format(vm.root_img, oldcow_img)],
            stdout=subprocess.PIPE)
        result = p.communicate()
        if result[0].strip() != old_dmdev:
            print >> sys.stderr, "ERROR: Cannot create snapshot device ({0} != {1})".format(
                    result[0].strip(), old_dmdev)
            exit(1)


    print "INFO: Reverting template changes done at {0}".format(oldcow_time_str)
    if not options.force:
        prompt = raw_input ("Do you want to proceed? [y/N] ")
        if not (prompt == "y" or prompt == "Y"):
            exit (0)
    
    p = subprocess.Popen(["/sbin/dmsetup", "table", old_dmdev], stdout=subprocess.PIPE)
    result = p.communicate()
    dm_table = result[0]
    dm_table_elements = dm_table.split(' ')
    if dm_table_elements[2] != 'snapshot':
        print >> sys.stderr, "ERROR: Unexpected device-mapper type ({0}). Template changes reverting already running".format(dm_table_elements[2])
        exit(1)

    dm_table_elements[2] = 'snapshot-merge'
    dm_table = ' '.join(dm_table_elements)
    subprocess.check_call(["/sbin/dmsetup", "reload", old_dmdev, "--table", dm_table])
    # Reload new table into LIVE slot
    subprocess.check_call(["/sbin/dmsetup", "suspend", old_dmdev])
    subprocess.check_call(["/sbin/dmsetup", "resume", old_dmdev])
    # Wait to snapshot merge completed
    while True:
        p = subprocess.Popen(["/sbin/dmsetup", "status", old_dmdev], stdout=subprocess.PIPE)
        result = p.communicate()
        status_details = result[0].split(' ')
        blocks_used = status_details[3].split('/')[0]
        if int(blocks_used) == int(status_details[4]):
            break
        print "\r-> Reverting template changes: {0} of {1} left".format(blocks_used, root_blocks),
        time.sleep(1)
    print "\r-> Reverting template changes: done                  ".format(blocks_used, root_blocks)

    dm_table_elements[2] = 'snapshot'
    dm_table = ' '.join(dm_table_elements)
    subprocess.check_call(["/sbin/dmsetup", "reload", old_dmdev, "--table", dm_table])
    # Reload new table into LIVE slot
    subprocess.check_call(["/sbin/dmsetup", "suspend", old_dmdev])
    subprocess.check_call(["/sbin/dmsetup", "resume", old_dmdev])

    subprocess.check_call(["/etc/xen/scripts/block-snapshot", "cleanup",
            "snapshot", old_dmdev])

    os.rename(oldcow_img, vm.rootcow_img)
    exit(0)


main()