#!/usr/bin/python2
# -*- encoding: utf8 -*-
#
# The Qubes OS Project, http://www.qubes-os.org
#
# Copyright (C) 2014 Matt McCutchen <matt@mattmccutchen.net>
# Copyright (C) 2015 Marek Marczykowski-Górecki <marmarek@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.
#
#

import os
import subprocess
import sys
import time
from qubes import qubes
from qubes import qubesutils


def is_dvm_up_to_date(tmpl, dvm_tmpl):
    dvm_savefile_path = os.path.join(dvm_tmpl.dir_path, "dvm-savefile")
    if not os.path.isfile(dvm_savefile_path):
        return False

    dvm_mtime = os.path.getmtime(dvm_savefile_path)
    root_mtime = os.path.getmtime(tmpl.root_img)
    if dvm_mtime < root_mtime:
        return False
    else:
        return True


def main():
    if len(sys.argv) != 2:
        print >> sys.stderr, 'Usage: qvm-trim-template TEMPLATEVM_NAME'
        sys.exit(1)

    tvm_name = sys.argv[1]

    qvm_collection = qubes.QubesVmCollection()
    qvm_collection.lock_db_for_writing()
    qvm_collection.load()

    tvm = qvm_collection.get_vm_by_name(tvm_name)
    if tvm is None:
        print >> sys.stderr, 'VM \'{}\' does not exists'.format(tvm_name)
        sys.exit(1)

    if not tvm.is_template():
        print >> sys.stderr, '{} is not template'.format(tvm_name)
        sys.exit(1)

    if tvm.is_running():
        print >> sys.stderr, 'Please stop the TemplateVM first.'
        sys.exit(1)

    outdated_children = [c for c in qvm_collection.get_vms_based_on(tvm.qid) if
                         c.is_outdated()]
    if outdated_children:
        print >> sys.stderr, 'Please stop (or restart) the following outdated VMs based on the template first:\n%s' % ', '.join(

            c.name for c in outdated_children)
        sys.exit(1)

    rootcow_old = tvm.rootcow_img + '.old'
    print 'Disk usage before:'
    subprocess.check_call(['du', tvm.root_img] + (
        ['--total', rootcow_old] if os.path.exists(rootcow_old) else []))

    # root-cow.img.old is likely to be invalid once we trim root.img, so go ahead and delete it.
    # (Note, root-cow.img should be logically empty because the TemplateVM is not running.)
    if os.path.exists(rootcow_old):
        os.remove(rootcow_old)

    dvm_tmpl = qvm_collection.get_vm_by_name(tvm.name + '-dvm')
    if dvm_tmpl is None:
        touch_dvm_savefile = False
    else:
        touch_dvm_savefile = is_dvm_up_to_date(tvm, dvm_tmpl)

    print >> sys.stderr, "Creating temporary VM..."
    fstrim_vm = qvm_collection.add_new_vm("QubesAppVm",
                                          template=tvm,
                                          name="{}-fstrim".format(tvm_name),
                                          netvm=None,
    )
    if not fstrim_vm:
        print >> sys.stderr, "ERROR: Failed to create new VM"
        sys.exit(1)

    fstrim_vm.create_on_disk()
    fstrim_vm.start(start_guid=False, verbose=True)

    print >> sys.stderr, "Performing fstrim now..."
    fstrim_process = fstrim_vm.run("/bin/sh", user="root", passio_popen=True,
                                   gui=False)
    fstrim_process.stdin.write('''
    until [ -r /dev/xvdi ]; do
      sleep 1
    done
    mkdir /tmp/root
    mount -o ro /dev/xvdi /tmp/root
    fstrim -v /tmp/root
    poweroff
    ''')
    fstrim_process.stdin.close()

    qubesutils.block_attach(qvm_collection, fstrim_vm,
                            {
                                'vm': 'dom0',
                                'device': tvm.root_img,
                                'mode': 'w',
                            },
                            mode='w',
                            frontend='xvdi')

    # At this point, the trim should run and the vm should shut down by itself and
    # detach the block device.

    fstrim_process.wait()
    print >> sys.stderr, "fstrim done, cleaning up..."

    while fstrim_vm.is_running():
        time.sleep(1)

    fstrim_vm.remove_from_disk()
    qvm_collection.pop(fstrim_vm.qid)

    # if DispVM template was up to date, keep that state
    if touch_dvm_savefile:
        dvm_savefile_path = os.path.join(dvm_tmpl.dir_path, 'dvm-savefile')
        os.utime(dvm_savefile_path, None)
        # If this is default DispVM, make sure that tmpfs copy of the file
        # (if enabled) also has mtime updated
        if os.stat('/var/lib/qubes/dvmdata/default-savefile').st_ino == \
                os.stat(dvm_savefile_path).st_ino:
            os.utime('/var/run/qubes/current-savefile', None)

    # do not save, all changes to qubes.xml should be reversed
    qvm_collection.unlock_db()

    print 'Disk usage after:'
    subprocess.check_call(['du', tvm.root_img])


if __name__ == "__main__":
    main()