# Copyright (c) 2016-2017 VMware, Inc.
# All rights reserved. -- VMware Confidential
"""
Integration of PostgreSQL with the B2B patching framework.

This follows the architecture specification defined here:
https://wiki.eng.vmware.com/VSphere2016/vSphere2016Upgrade/Inplace/ArchitectureSpec
As well the patching framework and hooks defined here:
https://wiki.eng.vmware.com/VSphere2016/vSphere2016Upgrade/Inplace/Patch_Extensibility
"""
import sys
import os
import json
import platform
import logging
import subprocess

from patch_specs import (DiscoveryResult, ValidationResult, Question,
                         Mismatch, Requirements, PatchInfo, RequirementsResult)
from extensions import extend, Hook
from l10n import msgMetadata as _T, localizedString as _
import os_utils
import rpm_utils
import vcsa_utils
from patch_errors import UserError
from reporting import getProgressReporter

logger = logging.getLogger(__name__)
MY_PAYLOAD_DIR = os.path.dirname(__file__)

# Set of utility routines
def getDbType():
   if os.environ.get('VMWARE_CFG_DIR') is None:
      logger.info("VMWARE_CFG_DIR is not defined.")
      return None
   VC_DB_TYPE = os.path.join(os.environ['VMWARE_CFG_DIR'], 'db.type')
   if not os.path.exists(VC_DB_TYPE):
      return None
   with open(VC_DB_TYPE) as fp:
      dbtype = fp.read()
   dbtype = dbtype.rstrip()
   return dbtype

def getPgVersion():
   ''' Get current PostgreSQL version of PGDATA'''
   if os.environ.get('VMWARE_POSTGRES_DATA') is None:
      logger.info("VMWARE_POSTGRES_DATA is not defined.")
      return None
   PG_VERSION_FILE = os.path.join(os.environ['VMWARE_POSTGRES_DATA'],
                                  'PG_VERSION')
   if not os.path.exists(PG_VERSION_FILE):
      return None
   with open(PG_VERSION_FILE) as fp:
      version = fp.read()
   version = version.rstrip()
   return version

def updatePgHealthStatusPath():
   '''Previous versions of postgres might be reporting health file at different
      locations, reconcile the health status file path configured in conf files
      generated at first boot time (e.g. health status worker conf, support
      manifest etc.) with pg vmon config specification of the health file (that
      is laid down through pg rpm).'''
   cfg_dir = os.environ['VMWARE_CFG_DIR']
   # conf file to read health status path from.
   pg_vmon_conf = os.path.join(cfg_dir, 'vmware-vmon', 'svcCfgfiles',
                               'vmware-vpostgres.json')
   # conf files to write health status path to.
   # (pg data dir is hardcoded since pg env vars are not available on VC's
   # super-main branch in patching sdk).
   pg_health_conf = os.path.join(os.environ['VMWARE_POSTGRES_DATA'],
                                 'health_status_worker.conf')
   pg_sup_man_conf = os.path.join(cfg_dir, 'vm-support',
                                  'vpostgres-support-noarch.mfx')
   with open(pg_vmon_conf, 'r') as f:
      vmon_vpg_cfg = json.load(f)
      if 'ApiHealthFile' in vmon_vpg_cfg:
         health_file = vmon_vpg_cfg['ApiHealthFile']
         # In-place update health status path using text substitution.
         # [ conf_file_path, 'regex_match_pattern', 'replacement_text']
         update_specs = [[pg_health_conf,
                          '^health_status_worker.status_file_path.*',
                          "health_status_worker.status_file_path = '%s'"
                          % health_file],
                         [pg_sup_man_conf, '^copy.*health-status.xml.*',
                          'copy IGNORE_MISSING %s' % health_file]]
         for (conf_file, pattern, substitution) in update_specs:
            update_command = ['sed', '-i', '-e',
                              's#%s#%s#g' % (pattern, substitution), conf_file]
            (stdout, stderr, res) = os_utils.executeCommand(update_command,
                                                            None)
            logger.info('%s update command run. stdout: %s, stderr: %s.' %
                        (conf_file, stdout, stderr))
            if res != 0:
               msg = 'PostgreSQL failed to run %s update command.' % conf_file
               raise UserError(_(_T('vpostgres.patch.healthconfupdate',
                                    msg)))


# Hook declarations and extensions
@extend(Hook.Discovery)
def discover(ctx): #pylint: disable=W0613
   '''Determine if a service wants to take part in the patching process.
      DiscoveryResult discover(PatchContext sharedCtx) throw UserError
   '''

   # Do not take part in patching process on Windows
   if platform.system().lower() == "windows":
      logger.info("No patching of PostgreSQL on Windows.")
      return None

   # Check if service is part of the infrastructure upgrade before looking
   # at if the database type is embedded. It would be possible that the file
   # to check if the database is embedded does not exist.
   result = vcsa_utils.getComponentDiscoveryResult('vpostgres')
   if not result or not result.patchable:
      return None

   # Check if the database type is external and do not take part
   # in patching process in this case.
   if getDbType() != 'embedded':
      logger.info("No patching of PostgreSQL if database type is not embedded.")
      return None

   # Establish the list of RPMs to ignore during installation. The only RPM
   # that needs to be taken into account here is VMware-Postgres-pg_archiver,
   # whose installation is done if the feature switch PG_WAL_ARCHIVE is
   # enabled.
   if not vcsa_utils.getTargetFSS('PG_WAL_ARCHIVE'):
      result.ignoreRpms.append('VMware-Postgres-pg_archiver')

   # It may be tempting to check what is the existing version of PostgreSQL
   # currently deployed at this stage, however it could be possible that
   # this process needs either a minor upgrade or a major upgrade. So delay
   # this decision-making in later stages to simplify the flow here.
   return result

@extend(Hook.SystemPrepare)
def prepare(ctx):
   '''void prepare(PatchContext sharedCtx) throw UserError'''

   # It is possible that the target build has the feature switch
   # PG_WAL_ARCHIVE disabled and that the source build has it enabled,
   # in which case make sure that in-place upgrades are able to complete
   # by forcibly removing the archiver package that is ignored at discovery
   # phase. Note that removeRpms is fail-safe, so even if the package is not
   # present because the source build has the feature switch disabled, this
   # will just log an error.
   if not vcsa_utils.getTargetFSS('PG_WAL_ARCHIVE'):
      rpm_utils.removeRpms(['VMware-Postgres-pg_archiver'])

@extend(Hook.Patch)
def patch(ctx):
   '''void patch(PatchContext sharedCtx) throw UserError'''
   progressReporter = getProgressReporter()

   progressReporter.updateProgress(0, _(_T("vpostgres.patch.begin",
                                           'PostgreSQL patch started.')))

   # Update health status file path as needed.
   updatePgHealthStatusPath()

   # Remove password expiration for OS user vpostgres to allow commands to
   # be executed in this account. Note that this is done before running the
   # upgrade commands as pg_upgrade and initdb are doing that.
   chage_command = ['/usr/bin/chage', '-M', '-1',
                    os.environ['VMWARE_POSTGRES_OS_ADMIN']]
   (stdout, stderr, res) = os_utils.executeCommand(chage_command, None)
   logger.info("chage command has run. stdout: %s, stderr: %s." % \
               (stdout, stderr))
   if res != 0:
      raise UserError(_(_T('vpostgres.patch.chage',
                           "Failed to disable password expiration for OS user %s." % \
                           (os.environ['VMWARE_POSTGRES_OS_ADMIN']))))

   # Check if an upgrade can be run here and do it. If nothing is necessary
   # the script would tell us that by itself. Note that after this command
   # has run PG_VERSION is updated as well if pg_upgrade has been invoked.
   # Note that this command is designed to be self-contained, so this file
   # should do nothing except calling this script. If only a minor upgrade
   # is needed, a couple of things can still happen but the binary called
   # here is smart enough to figure that out by itself.
   upgrade_command = os.path.join(os.environ['VMWARE_POSTGRES_BIN'],
                     'vmw_vpg_config/vmw_vpg_config.py')
   upgrade_command = [os.environ['VMWARE_PYTHON_BIN'], upgrade_command, '--CIS-mode',
                              '--action', 'upgrade_inplace']

   # Check if vmware-postgres-archiver service is expected to be functional
   # on this appliance, based on vMon service registration (which is simply
   # postgres-archiver rpm presence). In some appliances based on VC code base,
   # postgres-archiver service exists (e.g. vCenter Server Appliance),
   # in some it does not (e.g. VMC Gateway appliance does not install
   # postgres-archiver rpm).
   pg_archiver_exists = rpm_utils.isRpmInstalled('VMware-Postgres-pg_archiver')

   # Add switch equivalent to PG_WAL_ARCHIVE if enabled.
   if pg_archiver_exists and vcsa_utils.getTargetFSS('PG_WAL_ARCHIVE'):
      upgrade_command += ['--CIS-wal-archive']
   # Add switch equivalent to the feature switch PG_CHECKSUM.
   if vcsa_utils.getTargetFSS('PG_CHECKSUM'):
      upgrade_command += ['--CIS-checksum']
   (stdout, stderr, res) = os_utils.executeCommand(upgrade_command, None)
   logger.info("Upgrade command has run. stdout: %s, stderr: %s." % \
               (stdout, stderr))
   if res != 0:
      raise UserError(_(_T('vpostgres.patch.upgrade',
                           'PostgreSQL failed to run upgrade command.')))

   # Now that we are done with the upgrade, remove any upgrade-related
   # package. This is not considered as a failure, vCenter can live with
   # those packages still installed.
   rpm_list = []
   pg_ver = getPgVersion()
   if pg_ver >= '9.6':
      rpm_list += ['VMware-Postgres-upgrade-95']
   if pg_ver >= '9.5':
      rpm_list += ['VMware-Postgres-upgrade-94']
   if pg_ver >= '9.4':
      rpm_list += ['VMware-Postgres-upgrade-93']
   if len(rpm_list) != 0:
      rpm_utils.removeRpms(rpm_list)

   # Check that the partition dedicated to archives is present. If that's
   # not the case disable the archiver, this way it is still possible to
   # enable the service later on if the partition is added even after this
   # in-place upgrade. This only happens if the feature switch PG_WAL_ARCHIVE
   # is enabled.
   if (pg_archiver_exists and vcsa_utils.getTargetFSS('PG_WAL_ARCHIVE') and
      not os.path.ismount(os.environ['VMWARE_POSTGRES_MOUNT_ARCHIVE'])):
      vmon_command = os.path.join(os.environ['VMWARE_CIS_HOME'], 'vmware-vmon/vmon-cli')
      vmon_command = [vmon_command, '--starttype', 'DISABLED',
                      '--update', 'vmware-postgres-archiver']
      (stdout, stderr, res) = os_utils.executeCommand(vmon_command, None)
      logger.info("vmon-cli has disabled vmware-postgres-archiver. stdout: %s, stderr: %s." % \
                  (stdout, stderr))
      if res != 0:
         raise UserError(_(_T('vpostgres.patch.upgrade',
                              'PostgreSQL failed to run vmon-cli command.')))

   progressReporter.updateProgress(100,
        _(_T("vpostgres.patch.complete",
             'PostgreSQL patch completed.')))
