#!/usr/bin/python2
# -*- encoding: utf8 -*-

from qubes.qubes import QubesVmCollection
from qubes.qubes import QubesHost
from qubes.qubes import QubesException
from optparse import OptionParser
import subprocess 
import sys
import os
import re
import glob
import logging
import logging.handlers
from datetime import datetime
import time

LOG_FILENAME = '/var/log/qvm-update.log' 

def main():

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

    usage = "usage: %prog [options] <vm-name>"
    parser = OptionParser (usage)

    parser.add_option ("-S", "--shutdown", 
                       action="store_true",
                       dest="force_shutdown", 
                       default=[False],
                       help="shutdown system after completion")

    (options, args) = parser.parse_args ()

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

    nowstring = str(datetime.now()) 

    #now = datetime.now(); 
    my_logger = logging.getLogger('MyLogger')
    my_logger.setLevel(logging.INFO)
    handler = logging.handlers.RotatingFileHandler(LOG_FILENAME,maxBytes=200000,backupCount=5)
    my_logger.addHandler(handler)

    my_logger.info('')
    my_logger.info('qvm-update started:' + nowstring)

    print ''
    print "Checking for Dom0 updates"
    my_logger.info('Checking Dom0')

    cmd = "/usr/bin/sudo /usr/bin/qubes-dom0-update -y"
    my_logger.info(cmd)
    
    process = subprocess.Popen(cmd,shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    out, err = process.communicate()
    errcode = process.returncode

    dom0_updated = False 
    needs_restart = False 
    for line in out.split('\n'):
        my_logger.info(line)
        if re.search('^$',line): 
	    continue
        print '|' + line + '|'
	if re.search('^Updated:',line): 
	  dom0_updated = True
	if re.search('qubes-',line): 
	  needs_restart = True

    if errcode != 0:
	print "Error:" + str(errcode)
	print err
        my_logger.info(['Error:' + str(errcode)])
        my_logger.info(err)
    else:
        my_logger.info('Done.')
        print 'Done.'

    print '' 
    print "Checking all template VM's:"
    my_logger.info('Checking all template VM\'s:')

    vmnum = 0
    # for each template vm try to update it
    for vm in qvm_collection.values():
        if vm.is_template():
            vm_updated = False
            print ' ' 
            vmnum = vmnum + 1
            print str(vmnum) + ' ' + vm.name
	    cmd = "/usr/bin/qvm-run -a -p -u root " + vm.name + " \'yum -y update ; echo returned $? \'"
            #print cmd
            my_logger.info(str(vmnum) + ' ' + vm.name)
            my_logger.info(cmd)

    	    process = subprocess.Popen(cmd,shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
            out, err = process.communicate()
            errcode = process.returncode

            success = 0		
            for line in out.split('\n'):
                my_logger.info(line)
		if re.search('^out Loaded',line): 
		   continue
		if re.search('qubes-hooks',line): 
		   vm_updated = True
		   continue
		if re.search('^$',line): 
		   continue
                if re.search('returned 0',out):
                   print '  Done.'
                   success = 1
	           break 
                else:
                   print '|' + line + '|'

            if success != 1:
               if errcode != 1:
   	          print "Error:" + str(errcode)
   	          print err
                  my_logger.info('Errors:')
                  my_logger.info(err)
               else:
                   my_logger.info('Errors:')
                   for line in err.split('\n'):
                       my_logger.info(line)
                       if re.search('^Daemon',line): continue
	   	       if re.search('^Waiting for',line): continue
	   	       if re.search('^Connecting to',line): continue
	      	       if re.search('^Traceback',line): continue
	   	       if re.search('VM already stopped',line): continue
	   	       if re.search('main()',line): continue
	   	       if re.search('000QubesVm.py',line): continue
	   	       if re.search('vm.force_shutdown()',line): continue
	   	       if re.search('File "/usr/bin/qvm-run", line 253',line): continue
	   	       if re.search('^$',line): continue
	   	       print "Error: " + '|' + line + '|'
            vm.shutdown(force=True)
    
    # if we updated some templates then kick the dvm to save time loading first app
    if vm_updated == True :
        cmd = "sh -c 'echo bash -c exit | /usr/lib/qubes/qfile-daemon-dvm qubes.VMShell dom0 DEFAULT red'"
        my_logger.info([str(vmnum) + ' ' + vm.name])
        my_logger.info(cmd)
        process = subprocess.Popen(cmd,shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
        out, err = process.communicate()
        errcode = process.returncode
        print out

    print '' 
    nowstring = str(datetime.now()) 
    my_logger.info('qvm-update complete:' + nowstring)

    if options.force_shutdown == True:
        print "beginning shutdown sequence"
        nowstring = str(datetime.now()) 
        my_logger.info('shutting down system' + nowstring)
        shutdown_counter_max = 60
        shutdown_counter = 0 
        vms_list = []
        for vm in qvm_collection.values():
            if vm.is_running() and vm.type == 'AppVM': 
               vms_list.append(vm)
               print "initiating shutdown of AppVM: " + vm.name
               my_logger.info('initiating shutdown of AppVM: ' + vm.name)
               vm.shutdown(force=True)
        while len (vms_list):
            for vm in vms_list:
                my_logger.info('waiting on: ' + vm.name + "  " + str(shutdown_counter))
                if not vm.is_running():
                    vms_list.remove (vm)
                    my_logger.info('stopped: ' + vm.name)
                if shutdown_counter > shutdown_counter_max: 
                    print "forcing AppVM shutdown: " + vm.name 
                    my_logger.info('forcing AppVM shutdown: ' + vm.name)
                    if vm.is_running():
                        vm.force_shutdown()
            shutdown_counter += 1
            time.sleep (1)
        print "System shutdown"
        cmd = "shutdown -h now"
        my_logger.info('shutting down system' + nowstring)
        process = subprocess.Popen(cmd,shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
        out, err = process.communicate()
        errcode = process.returncode

main()