#!/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

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

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'])

'''
In some guests our udev magic doesn't work
'''
def check_bad_udev_guest(ve_uuid, tools_ver):
    print "CHECK ver" + 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
        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:
        return False

    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)
        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']
            if ve_db[ve_uuid]['type'] == 'Linux':
                new_ver = lin_guest_ver
            else:
                new_ver = win_guest_ver

            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 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())
        subprocess.call(['prlctl', 'exec', ve_uuid, 'mount', guest_dev, tmpdir.rstrip()])
        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"])
        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')
    for diskType in diskTypes:
        if diskType.getAttribute('type') != 'file' or diskType.getAttribute('device') != 'cdrom':
            continue
        diskNodes = diskType.childNodes
        iso_found = False
        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
    subprocess.call(['virsh', 'change-media', '--eject', '--force', uuid, iso_dev])
    return True


'''
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()
        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
    if ve_db[ve_uuid]['type'] == 'Linux':
        new_ver = lin_guest_ver
    else:
        new_ver = win_guest_ver
    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()
            update_exec(ve, ve_db, ve_uuid, True)

    # 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, ...]\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")

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] 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:
        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']
        if ve_db[ve_uuid]['type'] == 'Linux':
            new_ver = lin_guest_ver
        else:
            new_ver = win_guest_ver
        ve_db[ve_uuid]['last_update_to'] = new_ver
        ve_db[ve_uuid]['need_update'] = False

        processed_vms += 1
        if processed_vms == max_vms:
            log_debug("Reached max number of VMs to be processed per launch!")
            break

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