#!/opt/vmware/bin/python

import os
import logging
from cis.defaults import get_cis_data_dir
from cis.utils import run_command, create_dir, is_windows
import argparse
import json
import sys
from contextlib import contextmanager
import tempfile

if not is_windows():
    exit(1)
sys.path.append(os.path.join(os.environ['VMWARE_CIS_HOME'], 'vmafdd'))
import vmafd
from identity.vmkeystore import VmKeyStore
from cis.cmhelper import CISCmHelper

watchdogSubDirName = 'iiad'
watchdogMaintenanceModeName = 'iiad.maintenance-mode'


def getWatchdogSentinelFile():
    return os.path.join(get_cis_data_dir(), watchdogSubDirName,
                        watchdogMaintenanceModeName)


def createWatchdogSentinelFile():
    try:
        sentinelName = getWatchdogSentinelFile()
        with open(sentinelName, 'w') as sentinel:
            sentinel.write('stopped')
    except OSError:
        logging.error('Unable to create sentinel file %s.' % sentinelName)


def normPath(path):
    return os.path.normcase(os.path.normpath(path))


def getPreserveList():
    """
    Gets preservelist file path from VMWARE_UNINSTALL_PRESERVE env variable.
    Reads the preservelist file and returns a list of normalized file paths
    which need to preserved on uninstall.
    """
    preserveList = set()
    preserveListFile = os.environ.get('VMWARE_UNINSTALL_PRESERVE')
    if preserveListFile is not None:
        with open(preserveListFile, 'r') as fp:
            for filePath in fp:
                preserveList.add(normPath(os.path.expandvars(filePath)))
    return preserveList


def deleteDirectory(rootDir, preserveList):
    """
    Deletes all files and folder excluding ones in preserveList under
    rootDir. At the end will also delete the rootDir if empty.
    Note:- Fn doesn't halt on file/folder remove errors.

    Return:- A 2 elem tuple (success, childrenRemoved).
             success = True if no errors were encountered while deleting files
                       not in the preserve list.
             childrenRemoved = True if we were able to delete the rootDir.
    """
    if normPath(rootDir) in preserveList:
        return (True, False)

    success = childrenRemoved = True
    for fname in os.listdir(rootDir):
        fPath = normPath(os.path.join(rootDir, fname))
        if fPath in preserveList:
            childrenRemoved = False
        elif os.path.isdir(fPath):
            # Dir which is not in preserve list. Delete it.
            subDirSuccess, subDirChildRemoved = deleteDirectory(fPath,
                                                                preserveList)
            success = success if subDirSuccess else subDirSuccess
            childrenRemoved = childrenRemoved if subDirChildRemoved \
                                else subDirChildRemoved
        else:
            # Not in preserve list and is a file. Delete it.
            try:
                os.remove(fPath)
            except OSError as e:
                logging.warning('Unable to remove file %s. '
                                'Error: %s' % (fPath, e))
                childrenRemoved = success = False

    if childrenRemoved:
        # Delete the root directory since all children have been removed.
        try:
            os.rmdir(rootDir)
        except OSError as e:
            logging.warning('Unable to remove dir %s. '
                            'Error: %s' % (rootDir, e))
            childrenRemoved = success = False

    return (success, childrenRemoved)


@contextmanager
def tempinput(data):
    temp = tempfile.NamedTemporaryFile(delete=False)
    temp.write(data)
    temp.close()
    yield temp.name
    os.unlink(temp.name)


def getMachineAccountCertAndKey():
    """
    Get certificate and key for the machine account.
    """
    ks = VmKeyStore('VKS')
    ks.load('machine')
    cert = ks.get_certificate('machine')
    key = ks.get_key('machine')
    return (cert, key)


def deleteCMServiceEndpoints(ignoreErrors=True):
    """
    Delete all service endpoints from CM.
    """
    try:
        cert, key = getMachineAccountCertAndKey()
        vmafdClient = vmafd.client('localhost')
        with tempinput(key) as tempkeyname, tempinput(cert) as tempcertname:
            cisCmHelper = CISCmHelper(vmafdClient.GetCMLocation(),
                                      None, None,
                                      cert=tempcertname, key=tempkeyname)
            cisCmHelper.unregisterNodeFromCM(vmafdClient.GetPNID(),
                                             ignoreErrors=ignoreErrors)
    except Exception as e:
        logging.error('Failed to unregister node from CM. Error %s' % e)
        if not ignoreErrors:
            raise


def runUninstallFB(uninstallLogDir, fbStatusFile):
    """
    Run un-install section of the firstboot scripts
    """
    runFirstbootBin = os.path.join(os.environ['VMWARE_CIS_HOME'],
                                   'bin', 'run-firstboot-scripts.bat')
    uninstallFirstbootCmd = [runFirstbootBin, '--action', 'firstboot',
                             '--subaction', 'uninstall', '--logDirectory',
                             uninstallLogDir]
    try:
        # Run uninstall only on firstboots which got executed during install.
        # On successfull install we do this by not setting --fbWhiteList arg
        # of run-firstboot-scripts, which defaults to running all firsboot
        # scripts.
        # On failed install, we do this by setting --fbWhiteList arg to
        # stepsCompletedList + failedStep.
        with open(fbStatusFile, 'r') as fp:
            statusJsonData = json.load(fp)

        if statusJsonData['finalStatus'] == 'failure':
            fbWhiteList = statusJsonData['failedSteps']
            if statusJsonData['stepsCompletedList'].strip():
                fbWhiteList += ',' + statusJsonData['stepsCompletedList']
            uninstallFirstbootCmd.extend(['--fbWhiteList', fbWhiteList])
    except Exception as e:
        logging.warning('Error loading fbStatus.json file: %s' % e)

    rc, stdout, stderr = run_command(uninstallFirstbootCmd)
    if rc != 0:
        logging.error('Un-install scripts failed. Return code: %s, '
                      'Stdout: %s, Stderr: %s' % (rc, stdout, stderr))
        return False
    else:
        logging.info('Un-install scripts completed successfully.')
        return True


def collectVCSupport(vcSupportBundleDir):
    """
    Collect vc-support and store the .tgz in %VMWARE_LOG_DIR%/uninstall
    """
    try:
        # Create vcSupportBundleDir if it doesn't exist.
        create_dir(vcSupportBundleDir)
    except OSError as e:
        logging.error('Collecting vc-support failed. Error: %s' % e)
        return False

    vcSupportBin = os.path.join(os.environ['VMWARE_CIS_HOME'],
                                'bin', 'vc-support.bat')
    vcSupportCmd = [vcSupportBin, '-q', '-z', '-w', vcSupportBundleDir]
    rc, stdout, stderr = run_command(vcSupportCmd)
    if rc != 0:
        logging.error('Collecting vc-support failed. Return code: %s, '
                      'Stdout: %s, Stderr: %s' % (rc, stdout, stderr))
        return False
    else:
        logging.info('Collecting vc-support succeeded and '
                     'stored at: %s' % stdout)
        return True


def clearProgramData(preserveList):
    """
    Delete ProgramData files.
    """
    success = True
    for key in ['VMWARE_CFG_DIR', 'VMWARE_DATA_DIR',
                'VMWARE_RUNTIME_DATA_DIR', 'VMWARE_LOG_DIR']:
        path = os.environ.get(key)
        if path is not None and os.path.isdir(path):
            dirClearSuccess, _ = deleteDirectory(path, preserveList)
            if not dirClearSuccess:
                success = False

    return success

if __name__ == '__main__':
    """
    Cmd line input operation = UNINSTALL(Default), ROLLBACK

    -- UNINSTALL case --
    Run firstboot with uninstall flag.
        On success delete files/folder except the preserved files/folders.
        On failure collect VC support.

    -- ROLLBACK case --
    Run firstboot with uninstall flag.
        On success collect VC support and delete files/folder except the
        preserved files/folders.
        On Failure collect VC support.

    VMWARE_UNINSTALL_PRESERVE env variable can be used to specify a file
    containing files and dir which need to be preserved on UNINSTALL operation.

    Exit code:-
        0 :- Success
        1 :- Uninstall Failed.
        2 :- Failed to cleanup files in ProgramData.
    """
    # Setup logging.
    uninstallLogDir = os.path.join(os.environ['VMWARE_LOG_DIR'], 'uninstall')
    logFormat = '%(asctime)-15s: %(message)s'
    uninstallLogFile = os.path.join(uninstallLogDir, 'uninstall.log')
    create_dir(uninstallLogDir)
    logging.basicConfig(format=logFormat,
                        level=logging.INFO,
                        filename=uninstallLogFile,
                        filemode='a')

    # Parse cmd line arguments.
    allowedOps = ['UNINSTALL', 'ROLLBACK']
    parser = argparse.ArgumentParser()
    parser.add_argument('--operation', default='UNINSTALL',
                        help='Operations (%s)' % ' ,'.join(allowedOps))
    parser.add_argument('--fbStatusFile',
                        help='Path the firstboot status file.',
                        default=os.path.join(os.environ['VMWARE_LOG_DIR'],
                                             'firstboot',
                                             'firstbootStatus.json'))
    parser.add_argument('--vcSupportBundleDir', default=uninstallLogDir,
        help='Directory in which to place the VC support bundle.')
    cmdOptions = parser.parse_args()
    if cmdOptions.operation not in allowedOps:
        logging.error('Invalid operation specified %s' % cmdOptions.operation)
        exit(1)

    # preserveList will hold the list of files which we want to preserve.
    # Add uninstall.log file to this list. We will delete this at the end if
    # cleanup succeded.
    preserveList = set([normPath(uninstallLogFile)])

    # Remove entries from cm corresponding to all services.
    deleteCMServiceEndpoints()

    createWatchdogSentinelFile()
    # Run uninstall FB script.
    runUninstallFBSucceeded = runUninstallFB(uninstallLogDir,
                                             cmdOptions.fbStatusFile)

    rc = 0
    try:
        if cmdOptions.operation == 'UNINSTALL':
            # The VMWARE_UNINSTALL_PRESERVE hook is used only
            # in case of uninstall.
            preserveList.update(getPreserveList())

            if runUninstallFBSucceeded:
                if clearProgramData(preserveList):
                    preserveList.remove(normPath(uninstallLogFile))
                    logging.shutdown()
                    deleteDirectory(os.environ['VMWARE_LOG_DIR'], preserveList)
                else:
                    rc = 2
            else:
                # Collect VC support
                collectVCSupport(cmdOptions.vcSupportBundleDir)
                rc = 1
        elif cmdOptions.operation == 'ROLLBACK':
            # We should always collect VC support on rollback.
            collectVCSupport(cmdOptions.vcSupportBundleDir)
            if runUninstallFBSucceeded:
                # Preserve uninstall log dir. It contains the VC support.
                preserveList.add(normPath(uninstallLogDir))
                if not clearProgramData(preserveList):
                    rc = 2
    except Exception as e:
        # Don't think we can log anything as logging is shutdown.
        # Just going to call print till we can decide on a better way
        print str(e)

    exit(rc)
