#!/usr/bin/env python
# Copyright 2016-2022 VMware, Inc.
# All rights reserved. -- VMware Confidential

# This is a script which will be launched by the software-packages plugin
# prior to starting the RPM install transaction.
import sys
import os
import json
import subprocess
from os.path import join, abspath, dirname

from contextlib import contextmanager
import tempfile

import logging
import shutil
from preinstall_rpm_util import checkAndInstallRpms
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
logger.addHandler(logging.StreamHandler())

scriptsDir = abspath(dirname(__file__))
UPDATE_CONF = "/etc/applmgmt/appliance/update.conf"

@contextmanager
def atomic_write(filepath, newMode=0o644):
    '''
    Atomically write to a file
    '''

    oldStat = None
    try:
        oldStat = os.stat(filepath)
    except:
        pass

    tmpFile = tempfile.NamedTemporaryFile(
        mode='w+', dir=os.path.dirname(filepath), delete=False)
    rename = True
    try:
        yield tmpFile
    except:
        rename = False
        raise
    finally:
        tmpFile.flush()
        os.fsync(tmpFile.fileno())
        tmpFile.close()
        if rename:
            if oldStat:
                os.chmod(tmpFile.name, oldStat.st_mode)
                os.chown(tmpFile.name, oldStat.st_uid, oldStat.st_gid)
            else:
                os.chmod(tmpFile.name, newMode)
            os.rename(tmpFile.name, filepath)
        else:
            os.unlink(tmpFile.name)


# Load rpm-manifest.json to be used in case of no staged config
DEPLOYMENT_ROOT = abspath(join(dirname(__file__), os.pardir))
with open(join(DEPLOYMENT_ROOT,'rpm-manifest.json')) as f:
    rpmManifest = json.load(f)

try:
    STAGEDIR = '/storage/core/software-packages'
    with open(STAGEDIR + '/staged-configuration.json') as f:
        swUpdateState = json.load(f)
except:
    swUpdateState = rpmManifest

# If leaf node update, no point in doing any calculations.
# Just stop leaf service and its dependancy services and exit.
if rpmManifest.get('header', {}).get('vmware_stop_services'):
   stopServices = rpmManifest.get('header', {}).get('vmware_stop_services')
   logger.debug("Stop services are predefine: %s", stopServices)
   sys.exit(subprocess.call(['/bin/service-control', '--stop'] + stopServices))

sys.path.append(join(
    dirname(__file__),
    "patches",
    "libs",
    "feature-state"))
from update_utils import (runCommand, getSourceFSS, isGateway)
from product_utils import (getUpdateProduct, isHlmUpdate)

sys.path.extend(os.environ['VMWARE_PYTHON_PATH'].split(':'))

# This code is to suppress logs from featureState.py printed to stdout in
# cis.defaults. They init feature state again based on their internal flag
# cis.defaults.is_initialized.
import featureState
featureState.init(False)
import cis.defaults
cis.defaults.is_initialized = True

from cis.svcsController import (get_services_status,
                                update_services_runstate)

# PR 1324322 - Workaround for upgrading kernel RPMs.
# This directory must exist when the kernel rpms are being installed
# in order for the grub menu.lst file to be updated.
try:
    os.makedirs("/var/log/YaST2/")
except:
    pass

def modifyPythonLinks():
    '''
    Modify the links to point to python3
    :return:
    '''
    commonScript = join(scriptsDir, "common-patch.sh")
    out, err, rc = runCommand(["/bin/bash", "--login", "-c", commonScript],
                              progress=True,
                              message="Running common-patch script.")
    logger.debug("common-patch script \n out=%s\nerror=%s\nrc=%d\n" %
                 (out, err, rc))
    if rc:
        logger.debug("Running common-patch.sh failed.")

# TODO: (VCSA-3645)code cleanup after the leaf node service fix.
def calculateServicesToStop():
    '''
    Calculate the list of services that will be stopped during upgrade process
    @return servicesToStop - List of services that will be stopped
    @return nonLeafServiceWillBeUpdated - if any non-leaf service will be
        updated
    '''
    leafServices = swUpdateState['header'].get('leaf_services', [])
    logger.debug('Leaf services from patch: %r', leafServices)
    pkgList = swUpdateState['files']
    leafServicesToStop = []
    nonLeafServiceWillBeUpdated = False
    statuses = get_services_status(None)
    # if FSS on and we are on gateway and product is being updated then
    # don't consider core-services as stopping candidate
    if getSourceFSS('GW_PLATFORM_UPDATE') and getUpdateProduct():
        statuses = get_services_status(None, False, False)
        if isHlmUpdate():
            addtional_list = ["vmware-vmon"]
            for key, val in get_services_status(addtional_list).items():
                statuses[key] = val

    logger.debug("Statuses %r", statuses)
    for name in pkgList:
        if pkgList[name]['type'] != 'rpm':
            continue
        name = name.lower()
        if name not in leafServices:
            logger.debug('Non-leaf rpm found: %r', name)
            nonLeafServiceWillBeUpdated = True
        else:
            logger.debug('Leaf rpm found: %r', name)
            if name in statuses:
                logger.debug('Status: %r', statuses[name])
                v = statuses[name]
                if isinstance(v, tuple):
                    v = v[0]
                if statuses[name][0] == 'RUNNING':
                    logger.debug('Adding to stop-list: %r', name)
                    leafServicesToStop.append(name)
            else:
                logger.debug('Service %s is unknown', name)
    runningNonLeafServices = []
    for k, v in list(statuses.items()):
        if isinstance(v, tuple):
            v = v[0]
        if v == 'RUNNING' and k not in leafServices:
            logger.debug('Non-leaf service %s is running.'
                         ' Adding to the non-leaf stop-list', k)
            runningNonLeafServices.append(k)
    servicesToStop = leafServicesToStop
    if nonLeafServiceWillBeUpdated:
        servicesToStop.extend(runningNonLeafServices)

    return servicesToStop, nonLeafServiceWillBeUpdated


def stopAffectedServices(servicesToStop):
    '''
    Stop affected services
    '''
    if not servicesToStop:
        return
    update_services_runstate(
        'stop', None, False, False, servicesToStop,
        ignore_err=False)
    if 'vmware-stsd' in servicesToStop:
        #This is needed as a workaround for bug: 2162943
        logger.debug("Explicity killing all vmware-stsd.launcher processes.")
        subprocess.call(['/usr/bin/killall', 'vmware-stsd.launcher'])


servicesToStop, nonLeafServiceWillBeUpdated = calculateServicesToStop()

stopAffectedServices(servicesToStop)
if nonLeafServiceWillBeUpdated:
    modifyPythonLinks()

try:
    from subprocess import DEVNULL  # Python 3.
except ImportError:
    DEVNULL = open(os.devnull, 'wb')

if 'applmgmt' in servicesToStop:
    shutil.copyfile("/etc/applmgmt/firstboot/applmgmt.properties",
                    "/var/vmware/applmgmt/applmgmt.properties.bkp")
    logger.info('applmgmt is started as a standalone process under systemctl for patching')
    helperScript = join(scriptsDir, "pre-install-helper.sh")
    out, err, rc = runCommand(["/bin/bash", "--login", "-c", helperScript])
    if rc:
        logger.debug("Running pre-install-helper.sh failed. error=%s\n" % err)

# PR 2408571. Temporarily disable root password expration check for useradd and
# co.
with open('/etc/vmware/.buildInfo') as f:
    data = f.read()
# For CLoud Gateway appliance only
if 'CLOUDVM_NAME:VMware-vCenter-Cloud-Gateway' in data:
    pamd = '/etc/pam.d/'
    filelist = ['useradd', 'usermod', 'userdel', 'groupdel',
                'groupadd', 'chage']
    content = '''
    auth      required  pam_rootok.so
    account   required  pam_rootok.so
    '''
    logger.info('Replacing pam.d files')
    for fname in filelist:
        old_file = pamd + fname
        new_file = pamd + fname + '.patch.backup'
        if not os.path.exists(new_file):
            # Safeguard against situation when the patch is restarted
            # We do not want to overwrite backup files with out temporary
            # content
            shutil.copy2(old_file, new_file)  # Preserve mode, mtime, time
            with atomic_write(pamd + fname, 0o600) as f:
                f.write(content)
#[2396705] Explicitly making /var/run as symbolic link
os.system("if [ -d '/var/run' ] && [ ! -L '/var/run' ] ; then" \
          " mv /var/run/* /run; rm -rf /var/run; ln -sf /run /var/run;" \
          " echo '/var/run updated';"
          "fi")

#[2571185] Backup JRE Cert store
backupJreCertsScript = join(scriptsDir, "backup-jre-certs.sh")
newFileMode = 0o555
os.chmod(backupJreCertsScript, newFileMode)
out, err, rc = runCommand(["/bin/bash", "--login", "-c", backupJreCertsScript],
                            progress=True,
                            message="Running jre-cert-backup script.")
logger.info("backup-jre-certs script \n out=%s\nerror=%s\nrc=%d\n" %
                (out, err, rc))
if rc:
    logger.error("Running backup-jre-certs failed.")

def checkversion(UPDATE_CONF):
    version = None
    with open(UPDATE_CONF, "r") as f:
        updateConfJson = json.load(f)
        version = updateConfJson["version"]
    target_ver = "7.0.2.00000"
    src_ver = version
    if(src_ver == target_ver):
        return False
    src_ver_lst = src_ver.split(".")
    src_ver_lst = [int(i) for i in src_ver_lst]
    target_ver_lst = target_ver.split(".")
    target_ver_lst = [int(i) for i in target_ver_lst]

    if(src_ver_lst[0]!=target_ver_lst[0] or src_ver_lst[1]!=target_ver_lst[1]):
        return False

    if(target_ver_lst[2] > src_ver_lst[2]):
        return True

    return False

stageDirForrpm =  sys.argv[1]
stage_dir = os.path.abspath(os.path.join(stageDirForrpm, os.pardir))
compare_version = checkversion(UPDATE_CONF)
if(compare_version):
    checkAndInstallRpms(stage_dir)
