#!/usr/bin/python2.6
#
# The Qubes OS Project, http://www.qubes-os.org
#
# Copyright (C) 2010  Rafal Wojtczuk  <rafal@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 xen.lowlevel.xs
import os
import os.path
import sys
import subprocess
import daemon
import time
import dbus
from qubes.qubes import QubesVmCollection
from qubes.qubes import QubesException
from qubes.qubes import QubesDaemonPidfile
from qubes.qmemman_client import QMemmanClient

filename_seq = 50
pen_cmd = '/usr/lib/qubes/qubes_pencmd'
disposable_domains_dict = {}
current_savefile = '/var/run/qubes/current_savefile'

notify_object = None


def get_next_filename_seq():
    global filename_seq
    filename_seq = filename_seq + 1
    return str(filename_seq)

def logproc(msg):
        f = file('/var/log/qubes/qfileexchgd', 'a')
        f.write(msg+'\n')
        f.close()

def get_req_node(domain_id):
    return '/local/domain/'+domain_id+'/device/qpen'

def get_name_node(domain_id):
    return '/local/domain/'+domain_id+'/name'

def only_in_first_list(l1, l2):
    ret=[]
    for i in l1:
        if not i in l2:
            ret.append(i)
    return ret


class WatchType:
    def __init__(self, fn, param):
        self.fn = fn
        self.param = param

class DomainState:
    def __init__(self, domain, dict):
        self.rcv_state = 'idle'
        self.send_state = 'idle'
        self.domain_id = domain
        self.domdict = dict
        self.send_seq = None
        self.rcv_seq = None
        self.waiting_sender = None
        self.allowed_dest = None
        self.allowed_seq = None

    def killme(self):
        if not os.path.isfile('/etc/debug-dvm'):
            subprocess.call(['/usr/sbin/xm', 'destroy', self.domain_id])    
        
    def handle_request(self, request):
        req_ok = False
        if request is None:
            return
        tmp = request.split()
        rq = tmp[0]
        if len(tmp) > 1:
            vmname = tmp[1]
        else:
            vmname = None
        if len(tmp) > 2:
            transaction_seq = tmp[2]
        else:
            transaction_seq = '0'
        if rq == 'killme':
            self.killme()
            req_ok = True
        if rq == 'new' and self.send_state == 'idle':
            self.send_seq = get_next_filename_seq()
            retcode = subprocess.call([pen_cmd, 'new', self.domain_id, self.send_seq])
            logproc( 'Give domain ' + self.domain_id + ' a clean pendrive, retcode= ' + str(retcode))
            if retcode == 0:
                self.send_state = 'has_clean_pendrive'
            req_ok = True
        if rq == 'send' and self.send_state == 'has_clean_pendrive' and vmname is not None:
            logproc( 'send from ' + self.domain_id + ' to ' + vmname)
            if self.handle_transfer(vmname, transaction_seq):
                self.send_state = 'idle'
            req_ok = True;
        if rq == 'umount' and self.rcv_state == 'has_loaded_pendrive':
            retcode = subprocess.call([pen_cmd, 'umount', self.domain_id, self.rcv_seq])
            if retcode == 0:
                self.rcv_state = 'idle'
                self.rcv_seq = None
            logproc( 'set state of ' + self.domain_id + ' loaded->idle retcode=' + str(retcode))
            req_ok = True
        if rq == 'umount' and self.rcv_state == 'waits_to_umount':
            req_ok = True
            retcode = subprocess.call([pen_cmd, 'umount', self.domain_id, self.rcv_seq])
            if retcode != 0:
                return
            assert(self.waiting_sender != None)
            self.rcv_state = 'idle'
            self.rcv_seq = None
            tmp = self.waiting_sender
            self.waiting_sender = None
            if tmp.send_state == 'has_clean_pendrive':
                if tmp.handle_transfer(self.name, tmp.delayed_transaction_seq):
                    tmp.send_state = 'idle'

        if not req_ok:
            logproc( 'request ' + request + ' not served due to nonmatching state')

    def ask_to_umount(self, vmname):
        q = 'VM ' + vmname + ' has already an incoming pendrive, and thus '
        q+= 'cannot accept another one. If you intend to unmount its current '
        q+= 'pendrive and retry this transfer, press Yes. '
        q+= 'Otherwise press No to fail this transfer.'
        retcode = subprocess.call(['/usr/bin/kdialog', '--yesno', q, '--title', 'Some additional action required'])
        if retcode == 0:
            return True
        else:
            return False

    def handle_transfer_regular(self, vmname, transaction_seq):
        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:
            logproc( 'Domain ' + vmname + ' does not exist ?')
            return False
        if not vm.is_running():
            logproc( 'Domain ' + vmname + ' is not running ?')
            return False
        target=self.domdict[str(vm.get_xid())]
        if target.rcv_state != 'idle':
            if self.ask_to_umount(vmname):
                target.rcv_state='waits_to_umount'
                target.waiting_sender=self
                self.delayed_transaction_seq=transaction_seq
            logproc( 'target domain ' + target.domain_id + ' is not idle, now ' + target.rcv_state)
            return False
        if self.allowed_seq is not None:
            if self.allowed_seq != transaction_seq or self.allowed_dest != target.name:
                logproc('sender ' + self.name + ' receiver ' + target.name + ' : allowed attributes mismatch, denied')
                return False
        else:
            transaction_seq = '0'
            retcode = subprocess.call(['/usr/bin/kdialog', '--yesno', 'Do you authorize pendrive transfer from ' + self.name + ' to ' + vmname + '?' , '--title', 'Security confirmation'])
            logproc('handle_transfer: kdialog retcode=' + str(retcode))
            if retcode != 0:
                return False
        target.rcv_state='has_loaded_pendrive'
        retcode = subprocess.call([pen_cmd, 'send', self.domain_id, target.domain_id, self.send_seq, transaction_seq])
        target.rcv_seq = self.send_seq
        self.send_seq = None
        logproc( 'set state of ' + target.domain_id + ' to has_loaded_pendrive, retcode=' + str(retcode))
        if self.allowed_seq is not None:
            self.killme()
        return True

    def handle_transfer_disposable(self, transaction_seq):

        qmemman_client = QMemmanClient()
        if not qmemman_client.request_memory(400*1024*1024):
            qmemman_client.close()
            errmsg = 'Not enough memory to create DVM. '
            errmsg +='Terminate some appVM and retry.'
            subprocess.call(['/usr/bin/kdialog', '--sorry', errmsg])
            return False

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

        vm = qvm_collection.get_vm_by_name(self.name)
        if vm is None:
            logproc( 'Domain ' + vmname + ' does not exist ?')
            qvm_collection.unlock_db()
            qmemman_client.close()
            return False
        retcode = subprocess.call(['/usr/lib/qubes/qubes_restore',
            current_savefile,
            '-c', vm.label.color,
            '-i', vm.label.icon,
            '-l', str(vm.label.index)])
        qmemman_client.close()
        if retcode != 0:
            subprocess.call(['/usr/bin/kdialog', '--sorry', 'DisposableVM creation failed, see qubes_restore.log'])
            qvm_collection.unlock_db()
            return False
        f = open('/var/run/qubes/dispVM_xid', 'r');
        disp_xid = f.readline().rstrip('\n')
        disp_name = f.readline().rstrip('\n')
        disptempl = f.readline().rstrip('\n')
        f.close()
        vm_disptempl = qvm_collection.get_vm_by_name(disptempl);
        if vm_disptempl is None:
            logproc( 'Domain ' + disptempl + ' does not exist ?')
            qvm_collection.unlock_db()
            return False
        qvm_collection.add_new_disposablevm(disp_name, vm_disptempl.template_vm, label=vm.label)
        qvm_collection.save()
        qvm_collection.unlock_db()

        dispdom = DomainState(disp_xid, self.domdict)
        disposable_domains_dict[disp_xid] = dispdom
        retcode = subprocess.call([pen_cmd, 'send', self.domain_id, disp_xid, self.send_seq, transaction_seq])
        dispdom.rcv_seq = self.send_seq
        dispdom.rcv_state = 'has_loaded_pendrive'
        dispdom.allowed_dest = self.name
        dispdom.allowed_seq = transaction_seq 
        self.send_seq = None
        logproc( 'sent pendrive to disposable xid ' + disp_xid)
        return True

    def dvm_setup_ok(self):
        dvmdata_dir = '/var/lib/qubes/dvmdata/'
        if not os.path.isfile(current_savefile):
            return False
        if not os.path.isfile(dvmdata_dir+'default_savefile') or not os.path.isfile(dvmdata_dir+'savefile_root'):
            return False
        dvm_mtime = os.stat(current_savefile).st_mtime
        root_mtime = os.stat(dvmdata_dir+'savefile_root').st_mtime
        if dvm_mtime < root_mtime:
            return False       
        return True

    def tray_notify(self, str, timeout = 3000):
        notify_object.Notify("Qubes", 0, "red", "Qubes", str, [], [], timeout, dbus_interface="org.freedesktop.Notifications")

    def tray_notify_error(self, str, timeout = 3000):
        notify_object.Notify("Qubes", 0, "dialog-error", "Qubes", str, [], [], timeout, dbus_interface="org.freedesktop.Notifications")

    def handle_transfer(self, vmname, transaction_seq):
        if vmname != 'disposable':
            return self.handle_transfer_regular(vmname, transaction_seq)
        if not self.dvm_setup_ok():
            self.tray_notify("Updating DisposableVM savefile, please wait")
            if os.system("qvm-create-default-dvm --default-template --default-script") != 0:
                self.tray_notify_error("DVM savefile creation failed")
                return False 
        return self.handle_transfer_disposable(transaction_seq)

class XS_Watcher:
    def __init__(self):
        self.handle = xen.lowlevel.xs.xs()
        self.handle.watch('/vm', WatchType(XS_Watcher.dom_list_change, None))
        self.domdict = {}

    def dom_list_change(self, param):
        curr = self.handle.ls('', '/local/domain')
        if curr == None:
            return
        for i in only_in_first_list(curr, self.domdict.keys()):
            if disposable_domains_dict.has_key(i):
                newdom = disposable_domains_dict[i]
            else:
                newdom = DomainState(i, self.domdict)
                newdom.name = ''
            self.domdict[i] = newdom
            newdom.watch_token = WatchType(XS_Watcher.request, newdom)
            newdom.watch_name = WatchType(XS_Watcher.namechange, newdom)
            self.handle.watch(get_req_node(i), newdom.watch_token)
            self.handle.watch(get_name_node(i), newdom.watch_name)
            logproc( 'added domain ' + i)
        for i in only_in_first_list(self.domdict.keys(), curr):
            if disposable_domains_dict.has_key(i):
                self.remove_disposable_from_qdb(self.domdict[i].name)
                disposable_domains_dict.pop(i)
            self.handle.unwatch(get_req_node(i), self.domdict[i].watch_token)
            self.handle.unwatch(get_name_node(i), self.domdict[i].watch_name)
            self.domdict.pop(i)
            logproc( 'removed domain ' + i)

    def request(self, domain_param):
        ret = self.handle.read('', get_req_node(domain_param.domain_id))
        domain_param.handle_request(ret)

    def namechange(self, domain_param):
        ret = self.handle.read('', get_name_node(domain_param.domain_id))
	if ret!= '' and ret!=None:
		domain_param.name = ret
                logproc( 'Name for domain xid ' + domain_param.domain_id + ' is ' + ret )

    def remove_disposable_from_qdb(self, name):
        qvm_collection = QubesVmCollection()
        qvm_collection.lock_db_for_writing()
        qvm_collection.load()

        vm = qvm_collection.get_vm_by_name(name)
        if vm is None:
            logproc( 'remove_disposable_from_qdb: Domain ' + name + ' does not exist ?')
            qvm_collection.unlock_db()
            return False
        qvm_collection.pop(vm.qid)
        qvm_collection.save()
        qvm_collection.unlock_db()
        
    def watch_loop(self):
        global notify_object
        notify_object = dbus.SessionBus().get_object("org.freedesktop.Notifications", "/org/freedesktop/Notifications")

        sys.stderr = file('/var/log/qubes/qfileexchgd.errors', 'a')
        while True:
            result = self.handle.read_watch()
            token = result[1]
            token.fn(self, token.param)

def main():

    lock = QubesDaemonPidfile ("qfileexchgd")
    if lock.pidfile_exists():
        if lock.pidfile_is_stale():
            lock.remove_pidfile()
            print "Removed stale pidfile (has the previous daemon instance crashed?)."
        else:
            exit (0)


    context = daemon.DaemonContext(
        working_directory = "/var/run/qubes",
        pidfile = lock)
    with context:
        XS_Watcher().watch_loop()

main()