#!/usr/bin/python

#
# Copyright (c) 2017 Parallels International GmbH
#

import prlsdkapi
from prlsdkapi import consts as pc
import subprocess
import syslog
import time
import json
import sys
from string import ascii_uppercase
import libvirt
from xml.dom import minidom
import os
import re

prlsdkapi.init_server_sdk()
_server = prlsdkapi.Server()
_server.login_local().wait()

flags_running = [pc.VMS_STARTING, pc.VMS_RUNNING, pc.VMS_SUSPENDING, pc.VMS_SNAPSHOTING, pc.VMS_RESETTING, pc.VMS_PAUSING, pc.VMS_CONTINUING, pc.VMS_MOUNTED]

VE_DB = '/var/lib/vz-guest-tools-updater.json'
SUPPORTED_GUESTS = ['Linux', 'Windows']
DEBUG = False
GUEST_ISO_LIN = '/usr/share/vz-guest-tools/vz-guest-tools-lin.iso'
GUEST_MNT_POINT = '/tmp/vz-guest-tools-lin'

def log_debug(msg):
    if DEBUG:
        syslog.syslog(msg)

log_debug('Tools autoupdate started')

# We have to manually compare version of installed guest tools with version of guest tools package
# until we fix PSBM-60158
win_guest_ver = subprocess.check_output(['rpm', '-q', '--qf', '%{VERSION}-%{RELEASE}', 'vz-guest-tools-win'])
lin_guest_ver = subprocess.check_output(['rpm', '-q', '--qf', '%{VERSION}-%{RELEASE}', 'vz-guest-tools-lin'])

'''
Should we use version of Lin or Win guest tools?
'''
def get_actual_tools_ver(ve_type):
    if ve_type == 'Linux':
        return lin_guest_ver
    else:
        return win_guest_ver

'''
In some guests our udev magic doesn't work
'''
def check_bad_udev_guest(ve_uuid, tools_ver):
    log_debug("'Bad Udev' check for version " + tools_ver)
    if tools_ver.startswith("2.5"):
        return True

    FNULL = open(os.devnull, 'w')
    try:
        check_centos = subprocess.check_output(['prlctl', 'exec', ve_uuid, 'cat /etc/centos-release'], stderr=FNULL, close_fds=True)
        if "CentOS release 6" in check_centos:
            log_debug("CentOS 6 detected")
            return True
    except:
        log_debug("Centos check failed")
        pass

    try:
        check_debian = subprocess.check_output(['prlctl', 'exec', ve_uuid, 'cat /etc/debian_version'], stderr=FNULL, close_fds=True)
        if check_debian.startswith("7"):
            log_debug("Debian 7 detected")
            return True
    except:
        log_debug("Debian check failed")
        pass

    return False

'''
'Heavy' function - analyze all VMs and decide which of them should be updated
'''
def analyze_ves():
    # Load cache of VMs processed during the last launch
    try:
        ve_db = json.load(open(VE_DB))
    except:
        ve_db = {}

    # Get list of all VMs
    ves = _server.get_vm_list_ex(nFlags=pc.PVTF_VM).wait()

    processed_ves = []
    for ve in ves:
        ve_uuid = ve.get_uuid()
        log_debug("Processing " + ve_uuid)
        if ve_uuid not in ve_db:
            ve_db[ve_uuid] = {}
        processed_ves.append(ve_uuid)
        if ve.is_template():
            log_debug('Template, skipping')
            ve_db[ve_uuid]['need_update'] = False
            continue
        # 'type' can be missing in the db if autoupdate was previously disabled
        # so we didn't populate the db with all data
        if ve_uuid in ve_db and 'type' in ve_db[ve_uuid]:
            ve_type = ve_db[ve_uuid]['type']
            if ve_type not in SUPPORTED_GUESTS:
                continue
            if 'current' not in ve_db[ve_uuid]:
                ve_db[ve_uuid]['current'] = ''
            current_ver = ve_db[ve_uuid]['current']
            new_ver = get_actual_tools_ver(ve_db[ve_uuid]['type'])

            if current_ver == new_ver:
                if 'last_update_from' in ve_db[ve_uuid]:
                    syslog.syslog('{"name": "ToolsAutoUpdateResult", "actors": {"VEs": [{"%s"}]}, "from_version": "%s", "to_version": "%s", "result": "success"}' \
                            % (ve_uuid, ve_db[ve_uuid]['last_update_from'], ve_db[ve_uuid]['last_update_to']))
                ve_db[ve_uuid]['need_update'] = False
                ve_db[ve_uuid]['failed'] = ''
                continue
        else:
            ve_db[ve_uuid] = {}

        # Either we don't have VM in db or we want to refresh info about it
        conf = ve.get_config()

        ve_db[ve_uuid]['type'] = prlsdkapi.call_sdk_function('PrlApi_GuestToString', conf.get_os_type())

        vm_info = ve.get_vm_info()
        if vm_info.get_state() not in flags_running:
            log_debug('VM is not running, skipping')
            ve_db[ve_uuid]['need_update'] = False
            continue

        r = ve.get_tools_state().wait()
        if r.get_params_count() == 0:
            log_debug('Tools are not installed, nothing to update')
            ve_db[ve_uuid]['need_update'] = False
            continue

        tools_info = r.get_param()
        tools_ver = tools_info.get_version()
        ve_db[ve_uuid]['current'] = tools_ver

        guest_os_name = ve_db[ve_uuid]['type']
        if not conf.is_tools_auto_update_enabled():
            ve_db[ve_uuid]['need_update'] = False
        elif guest_os_name == "Linux":
            ve_db[ve_uuid]['need_update'] = (tools_ver != lin_guest_ver)
        elif guest_os_name == "Windows":
            ve_db[ve_uuid]['need_update'] = (tools_ver != win_guest_ver)
        else:
            ve_db[ve_uuid]['need_update'] = False

        if 'last_update_from' in ve_db[ve_uuid]:
            if ve_db[ve_uuid]['last_update_from'] == current_ver \
                    and ve_db[ve_uuid]['last_update_to'] == new_ver:
                if 'failed' in ve_db[ve_uuid] and ve_db[ve_uuid]['failed'] == 'once':
                    ve_db[ve_uuid]['failed'] = 'failed'
                    syslog.syslog('{"name": "ToolsAutoUpdateResult", "actors": {"VEs": [{"%s"}]}, "from_version": "%s", "to_version": "%s", "result": "failed"}' \
                            % (ve_uuid, ve_db[ve_uuid]['last_update_from'], ve_db[ve_uuid]['last_update_to']))
                    ve_db[ve_uuid]['need_update'] = False
                    continue
                else:
                    syslog.syslog('{"name": "ToolsAutoUpdateResult", "actors": {"VEs": [{"%s"}]}, "from_version": "%s", "to_version": "%s", "result": "failedOnce"}' \
                            % (ve_uuid, ve_db[ve_uuid]['last_update_from'], ve_db[ve_uuid]['last_update_to']))
                    ve_db[ve_uuid]['failed'] = 'once'
        else:
            # Hack for CentOS 6.x - our udev rule doesn't work there for some reason so let's force
            # update thorugh exec
            if ve_db[ve_uuid]['type'] == 'Linux':
                if check_bad_udev_guest(ve_uuid, tools_ver):
                    ve_db[ve_uuid]['failed'] = 'once'

    # Clean VE_DB - drop entries that don't exist anymore
    to_drop = []
    for ve in ve_db:
        if not ve in processed_ves:
            to_drop.append(ve)
    for ve in to_drop:
        ve_db.pop(ve, None)

    open(VE_DB, 'w').write(json.dumps(ve_db, indent = 2))

'''
Read the database and return list of VMs to be updated
'''
def get_ves_to_update():
    try:
        ve_db = json.load(open(VE_DB))
    except:
        log_debug("Failed to load json")
        ve_db = {}

    ve_list = []
    for ve_uuid in ve_db:
        if 'need_update' in ve_db[ve_uuid]:
            if ve_db[ve_uuid]['need_update']:
                log_debug("Will auto process: " + ve_uuid)
                ve_list.append(ve_uuid)

    return ve_list

'''
Get maximum number of VMs that can be processed at once
'''
def get_max_vms_per_time():
    # Limit number of VMs that can be processed at once
    try:
        conf = json.load(open('/etc/vz/tools-update.conf'))
        max_vms = conf['MaxVMs']
    except:
        # Just a hardcoded value - we use it in our config by default
        max_vms = 5
    return max_vms

'''
Update install-tools script inside the guest before connecting CD with guest tools
In old guest tools, that script could be killed by timeout.
'''
def update_udev_script(ve_uuid):
    log_debug("Updating install-tools script")
    if not os.path.exists(GUEST_MNT_POINT):
        os.mkdir(GUEST_MNT_POINT)
    else:
        subprocess.call(['umount', GUEST_MNT_POINT])

    if subprocess.call(['mount', '-o', 'loop', GUEST_ISO_LIN, GUEST_MNT_POINT]) != 0:
        log_debug("Failed to mount guest cd on host")
        return False
    try:
        newf = open(GUEST_MNT_POINT + "/install-tools")
        subprocess.call(['prlctl', 'exec', ve_uuid, "cat /dev/null > /usr/bin/install-tools.new"])
        for l in newf.readlines():
            l = re.escape(l).rstrip()
            subprocess.call(['prlctl', 'exec', ve_uuid, "echo " + l + " >> /usr/bin/install-tools.new"])
        subprocess.call(['prlctl', 'exec', ve_uuid, "cat /usr/bin/install-tools.new > /usr/bin/install-tools"])
        subprocess.call(['prlctl', 'exec', ve_uuid, "rm -f /usr/bin/install-tools.new"])
    except:
        log_debug("Smth went wrong during update")
        pass

    subprocess.call(['umount', GUEST_MNT_POINT])
    return True

'''
Update guest tools inside VM using direct execution of necessary actions
through 'prlctl exec'
'''
def update_exec(ve, ve_db, ve_uuid, linux_skip_run):
    guest_dev = ""
    if ve_db[ve_uuid]['type'] == 'Linux':
        log_debug("Looking for cd device in Lin guest")
        # Let guest some time to detect a new device
        time.sleep(5)
        proc = subprocess.Popen(['prlctl', 'exec', ve_uuid, 'blkid'], stdout=subprocess.PIPE)
        for line in proc.stdout:
            if "vz-tools-lin" in line:
                guest_dev = line.split(":")[0]
                break
        if not guest_dev:
            log_debug("Failed to detect guest cd")
            return False
        log_debug("Detected: " + guest_dev)
        tmpdir = subprocess.check_output(['prlctl', 'exec', ve_uuid, 'mktemp', '-d'])
        log_debug("Will mount to temp dir " + tmpdir.rstrip())
        if subprocess.call(['prlctl', 'exec', ve_uuid, 'mount', guest_dev, tmpdir.rstrip()]) != 0:
            log_debug("Failed to mount guest cd")
            return False
        log_debug("Copying install-tools...")
        subprocess.call(['prlctl', 'exec', ve_uuid, '/bin/cp', '-f', tmpdir.rstrip() + "/install-tools", "/usr/bin"])
        if not linux_skip_run:
            log_debug("Launching install-tools...")
            subprocess.call(['prlctl', 'exec', ve_uuid, "/usr/bin/install-tools"])
        else:
            if subprocess.call(['prlctl', 'exec', ve_uuid, 'umount', guest_dev]) != 0:
                log_debug("Failed to umount guest cd")
                return False
        log_debug("Done!")
    elif ve_db[ve_uuid]['type'] == 'Windows':
        log_debug("Looking for cd device in Win guest")
        # Iterate through all possible drive letters
        for drive in ascii_uppercase:
            proc = subprocess.Popen(['prlctl', 'exec', ve_uuid, 'vol', drive], stdout=subprocess.PIPE)
            for line in proc.stdout:
                if "vz-tools-win" in line:
                    guest_dev = drive
                    break
            if guest_dev:
                break
        if not guest_dev:
            return False
        log_debug("Detected: " + guest_dev)
        subprocess.call(['prlctl', 'exec', ve_uuid, guest_dev + ":/setup.exe"])

    return True

'''
Disconnect iso with guest tools from VM if connected.
If iso is already connected then 'installtools' does nothing, we have
no events inside VM guest so update is not triggered.

We have to use libvirt API here since our SDK "hides" guest tools CD
'''
def disconnect_tools_cd(conn, ve, ve_uuid):
    uuid = ve_uuid.replace("{", "").replace("}","")
    dom = conn.lookupByUUIDString(uuid)
    if dom == None:
        log_debug('Failed to find the domain ' + ve_uuid)
        return False

    raw_xml = dom.XMLDesc(0)
    xml = minidom.parseString(raw_xml)
    diskTypes = xml.getElementsByTagName('disk')
    iso_found = False
    for diskType in diskTypes:
        if diskType.getAttribute('type') != 'file' or diskType.getAttribute('device') != 'cdrom':
            continue
        diskNodes = diskType.childNodes
        for diskNode in diskNodes:
            if diskNode.nodeName != 'source' and not iso_found:
                continue
            if not iso_found:
                if not 'file' in diskNode.attributes.keys():
                    continue
                if not diskNode.getAttribute('file').endswith('vz-guest-tools-lin.iso'):
                    continue
                iso_found = True
                continue
            if diskNode.nodeName == 'target':
                iso_dev = diskNode.getAttribute('dev')
                break
        if iso_found:
            if not iso_dev:
                log_debug("Found connected guest tools CD, but failed to connect its parameters")
                return False
            break

    if not iso_found:
        # Guest tools cd is not mounted yet
        return True

    log_debug("Will disconnect " + iso_dev)
    # Dunno how to do this using Python API
    if subprocess.call(['virsh', 'change-media', '--eject', '--force', uuid, iso_dev]) != 0:
        log_debug("Failed to eject media using virsh")
        return False

    time.sleep(5)
    return True

'''
Check if tools in given VEs are up-to-date
'''
def check_update(ve_names):
   ves = _server.get_vm_list_ex(nFlags=pc.PVTF_VM).wait()
   for ve in ves:
       ve_uuid = ve.get_uuid()
       if ve_uuid not in ve_names and ve.get_name() not in ve_names:
            continue
       ve_uuid = ve.get_uuid()
       conf = ve.get_config()
       vm_info = ve.get_vm_info()
       ve_desc = ve_uuid + " (" + ve.get_name() + ")"
       if vm_info.get_state() not in flags_running:
           print('VM %s is not running, skipping' % ve_desc)
           continue
       ve_type = prlsdkapi.call_sdk_function('PrlApi_GuestToString', conf.get_os_type())
       if ve_type not in SUPPORTED_GUESTS:
           print('Unsupported guest typei in %s' % ve_desc)
           continue
       r = ve.get_tools_state().wait()
       if r.get_params_count() == 0:
           print('%s: Tools are not installed' % ve_desc)
           continue

       tools_info = r.get_param()
       current_ver = tools_info.get_version()
       new_ver = get_actual_tools_ver(ve_type)

       if not current_ver:
           print('%s: Tools are not installed' % ve_desc)
           continue
    
       if current_ver == new_ver:
           print('%s: Tools are up to date' % ve_desc)
       else:
           print('%s: Tools are outdated (%s vs %s)' % (ve_desc, current_ver, new_ver))
    

'''
Check if tools should be updated in a given VE and trigger update, if yes
'Force' argument forces tools update for VE with 'auto_update' set to false
Return True if update was triggered, False otherwise
'''
def process_ve(ve, ve_db, conn, force=False):
    ve_uuid = ve.get_uuid()
    log_debug('Processing %s ' % (ve_uuid))
    conf = ve.get_config()
    # Just for sure, check that user didn't turn off autoupdate
    if not force and not conf.is_tools_auto_update_enabled():
        log_debug('Autoupdate disabled, skipping')
        return False
    # Check that VM is running
    vm_info = ve.get_vm_info()
    if vm_info.get_state() not in flags_running:
        log_debug('VM is not running, skipping')
        return False

    # It is possible that we don't have this VM into the db
    # if VM was specified in command line explicitly
    # or have not 'type' set (e.g., if autoupdate was disabled during 'analyze' phase)
    if ve_uuid not in ve_db or 'type' not in ve_db[ve_uuid]:
        ve_db[ve_uuid] = {}
        ve_db[ve_uuid]['type'] = prlsdkapi.call_sdk_function('PrlApi_GuestToString', conf.get_os_type())
        if ve_db[ve_uuid]['type'] not in SUPPORTED_GUESTS:
            log_debug('Unsupported guest type')
            return False

        r = ve.get_tools_state().wait()
        if r.get_params_count() == 0:
            log_debug('Tools are not installed, nothing to update')
            ve_db[ve_uuid]['need_update'] = False
            return False

        tools_info = r.get_param()
        tools_ver = tools_info.get_version()
        if not tools_ver:
            log_debug('Tools are not installed, nothing to update')
            ve_db[ve_uuid]['need_update'] = False
            return False
        ve_db[ve_uuid]['current'] = tools_ver

        # Hack for CentOS 6.x - our udev rule doesn't work there for some reason so let's force
        # update thorugh exec
        if check_bad_udev_guest(ve_uuid, tools_ver):
            ve_db[ve_uuid]['failed'] = 'once'

    log_debug("Update has been triggered!")
    # JSON record for the CEP collector
    new_ver = get_actual_tools_ver(ve_db[ve_uuid]['type'])
    syslog.syslog('{"name": "ToolsAutoUpdate", "actors": {"VEs": [{"%s"}]}, "start_time": "%s", "from_version": "%s", "to_version": "%s"}' \
                    % (ve.get_uuid(), time.time(), ve_db[ve_uuid]['current'], new_ver))

    # Check and fix missing install-tools file bug (u3 legacy, Linux guests only)
    if ve_db[ve_uuid]['type'] == 'Linux':
        if subprocess.call(['prlctl', 'exec', ve_uuid, 'ls', '/usr/bin/install-tools']) != 0:
            # File is missing - copy it from CD
            # TODO: Maybe unpack it from ISO image and copy directly instead of working inside guest?
            log_debug('install-tools is missing, copying it to guest')
            disconnect_tools_cd(conn, ve, ve_uuid)
            ve.install_tools()
            if not update_exec(ve, ve_db, ve_uuid, True):
                # No need to continue if smth went wrong.
                # However, indicate that we have tried to do smth with VM
                return True
        elif ve_db[ve_uuid]['current'].startswith("0.10-"):
            # Checkif we have to preliminary update install-tools script inside the guest
            pkg_release = ve_db[ve_uuid]['current'].replace("0.10-","")
            pkg_release = pkg_release.replace(".vz7","")
            if int(pkg_release) < 112:
                update_udev_script(ve_uuid)

    # Disconnect tools iso, if connected, otherwise 'installtools'
    # will have no effect
    disconnect_tools_cd(conn, ve, ve_uuid)

    # We need to call installtools in either case
    ve.install_tools()
    if 'failed' in ve_db[ve_uuid] and ve_db[ve_uuid]['failed'] != "":
        # If update was failed once, then likely installtools by itself is not enough
        # let's try to launch update through exec
        if ve_db[ve_uuid]['failed'] != 'once' and not force:
            log_debug('Unexpected value of "failed", refusing to trigger update')
            return False
        else:
            log_debug('Triggering update using exec')
            update_exec(ve, ve_db, ve_uuid, False)
            ve_db[ve_uuid]['failed'] = ""
    return True

def print_help():
    prog_name = sys.argv[0]
    print("usage: %s [-h|--help] [--analyze|vm1, vm2, ...|--get-state vm1, vm2, ...]\n" % prog_name)
    print("%s - a tool for automated update of guest tools inside Virtual Machines.\n" % prog_name)
    print("If launched with '--analyze' option, the tool analyzes the state of every VM ")
    print("  registered in the system with 'GuestTools autoupdate' parameter set to 'on' and chooses VMs ")
    print("  with outdated guest tools. This information is stored in %s file.\n" % VE_DB)
    print("If launched without any arguments, %s triggers guest tools update in VMs" % prog_name)
    print("  from the %s list. " % VE_DB)
    print("  Maximum number of VMs where update can be triggered during a single guest tools updater")
    print("  invocation is limited by 'MaxVMs' parameter in the /etc/vz/tools-update.conf configuration file.\n")
    print("Alternatively, one can explicitly specify names or UUIDs of VMs where guest tools update ")
    print("  should be triggered. In this case, update is launched in all specified VMs at once regardless")
    print("  of their parameters and version of installed tools.\n")
    print("--get-state option can be used to check if guest tools are up-to-date in given VMs\n")

if __name__ == '__main__':
    ve_names = []
    # For manually specified VEs, we will force update even if
    # auto_update parameter is set to False
    force_update = False
    if len(sys.argv) > 1:
        if sys.argv[1] == "--analyze":
            if "-d" in sys.argv or "--debug" in sys.argv:
                DEBUG = True
            analyze_ves()
            sys.exit(0)
        if sys.argv[1] == "--get-state":
            if len(sys.argv) <= 2:
                print_help()
                sys.exit(0)
            ve_names = sys.argv[2:]
            check_update(ve_names)
            sys.exit(0)
        if sys.argv[1] in ["-h", "--help"]:
            print_help()
            sys.exit(0)
        if sys.argv[1] in ["-d", "--debug"]:
            DEBUG = True
            if len(sys.argv) > 2:
                ve_names = sys.argv[2:]
            else:
                ve_names = get_ves_to_update()
        else:
            ve_names = sys.argv[1:]
        force_update = True
    else:
        ve_names = get_ves_to_update()

    if len(ve_names) == 0:
        sys.exit(0)
    ves = _server.get_vm_list_ex(nFlags=pc.PVTF_VM).wait()

    max_vms = get_max_vms_per_time()
    processed_vms = 0
    ve_db = json.load(open(VE_DB))

    conn = libvirt.open('qemu:///system')
    if conn == None:
        print("Failed to open connection to qemu:///system")
        sys.exit(1)

    for ve in ves:
        if processed_vms >= max_vms:
            log_debug("Reached max number of VMs to be processed per launch!")
            break

        ve_uuid = ve.get_uuid() 
        if ve_uuid not in ve_names and ve.get_name() not in ve_names:
            continue
        if not process_ve(ve, ve_db, conn, force_update):
            continue
        ve_db[ve_uuid]['last_update_from'] = ve_db[ve_uuid]['current']
        new_ver = get_actual_tools_ver(ve_db[ve_uuid]['type'])
        ve_db[ve_uuid]['last_update_to'] = new_ver
        ve_db[ve_uuid]['need_update'] = False

        processed_vms += 1

    conn.close()
    open(VE_DB, 'w').write(json.dumps(ve_db, indent = 2))
