# Copyright (c) 2019-2021 VMware, Inc.  All rights reserved. -- VMware Confidential
# coding: utf-8

"""
All patching payload for wcpsvc.yaml are added in this file, for more info of
how to add your patches, please refer to:
https://confluence.eng.vmware.com/display/WCP/Run+VC+B2B+patching+test+before+checking+in+your+change
"""

import logging
import yaml
import os
import subprocess
import shutil
from cis.utils import FileBuffer
from cis.tools import get_install_parameter

logger = logging.getLogger(__name__)


# Add the description of your changes to CHANGELOG
CHANGELOG = ["Persist newly added Auto-upgrade config values in"
             " wcpsvc.yaml after patch when they are changed locally.",
             "Check if a config paramater is present in the old wcpsvc.yaml"
             "before trying to preserve its values"
             ]

# If changed locally, wcpsvc.yaml will be renamed as wcpsvc.yaml.rpmsave
# upon installing wcpsvc.rpm, which brings a new wcpsvc.yaml file.
PRESERVED_WCPSVC_YAML = "/etc/vmware/wcp/wcpsvc.yaml.rpmsave"
WCPSVC_YAML = "/etc/vmware/wcp/wcpsvc.yaml"
SAVED_WCPSVC_YAML = "/etc/vmware/wcp/wcpsvc.yaml.save"


# Entries allowlisted in wcpsvc.yaml and their default value, which means
# even values are changed locally on VC, changed values of these allow
# listed entries will be pertained after patch.
# It is subject to be extended when developers want to add new entries in
# wcpsvc.yaml which need to be pertained when they're changed locally
CONFIG_KEYS_ALLOWLIST = {
   "logging": {"level": "debug", "maxsizemb": "10", "maxbackups": "5"},
   "clusterconfig": {"upgrade_precheck_timeout": "300",
                     "dhcp_duid_organization_id": "6876"},
   "hdcsconfig": {"password_rotate_days": "60"},
   "autoupgradeconfig": {"num_of_clusters_upgradable_in_parallel": "5",
                         "auto_upgrade_poll_interval": "300"},
   "tkg_repository_config": {
      "repository_url": "https://wp-content.vmware.com/v2/latest/lib.json"
   }
}

"""
Define your patching functions below
"""


def preservePreviousEntries(old_wcpsvc_yaml=PRESERVED_WCPSVC_YAML,
                            new_wcpsvc_yaml=WCPSVC_YAML):
   """
   Function to retain allow listed entries in old wcpsvc.yaml when they
   are changed locally.
   If there are changes made to wcpsvc.yaml file locally and new wcpsvc.rpm
   file contains changes to wcpsvc.yaml file. The original wcpsvc.yaml will
   be renamed to wcpsvc.yaml.rpmsave after rpm installation. All the
   allow-listed entries in old wcpsvc.yaml will be retained and carried over
   to new wcpsvc.yaml brought by installation of wcpsvc.rpm when they are
   changed locally.

   :param old_wcpsvc_yaml: Old wcpsvc.yaml which is renamed after wcpsvc.rpm
   installed.
   :type old_wcpsvc_yaml: str
   :param new_wcpsvc_yaml: New wcp version file brought by wcpsvc.rpm
   installation.
   :type new_wcpsvc_yaml: str
   :rtype: None
   """
   if not os.path.exists(old_wcpsvc_yaml):
      logger.info("wcpsvc.yaml file is not modified locally, "
                  "no modifications to carry over.")
      return

   wcpsvc_yaml_save = None
   wcpsvc_yaml_new = None
   with open(old_wcpsvc_yaml) as fp:
      wcpsvc_yaml_save = yaml.safe_load(fp)
   with open(new_wcpsvc_yaml) as fp:
      wcpsvc_yaml_new = yaml.safe_load(fp)

   # Read allowlisted entries in old wcpsvc.yaml and overwrite them to
   # corresponding entries in new wcpsvc.yaml after patch
   # In short term we only consider 2-level-entry patching in wcpsvc.yaml,
   # will expand the multi-level wcpsvc.yaml patching when the dependent
   # bug: https://bugzilla.eng.vmware.com/show_bug.cgi?id=2459425 is
   # resolved.
   for category in CONFIG_KEYS_ALLOWLIST.keys():
      for key in CONFIG_KEYS_ALLOWLIST[category]:
         logger.info("category: {0}, key: {1}".format(category, key))
         if (category not in wcpsvc_yaml_new or key
            not in wcpsvc_yaml_new[category]):
            logger.info("Category {0} and entry {1} doesn't exist in new"
                        "wcpsvc.yaml after patching. No need to carry it over"
                        .format(category, key))
            continue
         try:
            # The FileBuffer() usage to replace entries in wcpsvc.yaml
            # is in short term since the python module ruamel.yaml
            # is not packaged in photonOS, bug to track:
            # https://bugzilla.eng.vmware.com/show_bug.cgi?id=2459425
            # In long term, FileBuffer() will be replaced with
            # ruamel.yaml to be used to replace entries in wcpsvc.yaml
            # while pertaining all comments.
            fb = FileBuffer()
            fb.readFile(new_wcpsvc_yaml)
            old_entry = key + ": " + CONFIG_KEYS_ALLOWLIST[category][key]
            if (category not in wcpsvc_yaml_save or key
               not in wcpsvc_yaml_save[category]):
               logger.info("Category {0} and entry {1} doesn't exist in old "
                        "wcpsvc.yaml. Adding them to new wcpsvc.yaml"
                        .format(category, key))
               new_entry = key + ": " + CONFIG_KEYS_ALLOWLIST[category][key]
            else:
               logger.info("Carrying entry {0} with value {1} from old "
                        "wcpsvc.yaml to newly patched wcpsvc.yaml file."
                        .format(key, wcpsvc_yaml_save[category][key]))
               new_entry = key + ": " + str(wcpsvc_yaml_save[category][key])
            fb.findAndReplace(old_entry, new_entry)
            fb.writeFile(new_wcpsvc_yaml)
         except Exception as ex:
            logging.exception("Failed to update {0} in new wcpsvc.yaml: {1}"
                              .format(key, ex))
            raise

   logger.info("All allowlisted entries in old wcpsvc.yaml have been carried"
               " over to newly patched wcpsvc.yaml")

   # Remove the preserved wcpsvc.yaml.save after we extract & carry over
   # allowlisted entries from it. In case for the next patch, that file would
   # still exist and we would unnecessarily try to again extract and carry
   # over entries.
   subprocess.run(["rm", old_wcpsvc_yaml])

   logger.info("wcpsvc.yaml patching completed. Preserved file {0} deleted."
               .format(old_wcpsvc_yaml))


def substituteRhttpProxyPort(wcp_config_file_path=WCPSVC_YAML):
   """
   Stores rhttpproxy port set during firstboot in wcpsvc config file.
   """

   # Default rhttp proxy port number
   _RHTTPPROXY_PORT_DEFAULT = '443'

   # Key name of entry rhttpproxy in wcpsvc.yaml file
   _RHTTP_PROXY = 'rhttpproxy_port'
   # Default value of entry rhttpproxy in wcpsvc.yaml
   _RHTTP_PROXY_DEF_VAL = 'rhttpproxy.ext.port2'

   with open(wcp_config_file_path) as fp:
      wcpsvc_yaml = yaml.safe_load(fp)

   # If there is no rhttpproxy entry in installed wcpsvc.yaml, which means
   # rhttpproxy entry is not needed in wcpsvc.yaml. Then skip the patch.
   if _RHTTP_PROXY not in wcpsvc_yaml.keys():
      logging.info('No entry {0} required in {1}, skip the patch.'
                   .format(_RHTTP_PROXY, wcp_config_file_path))
      return

   # If there is already value set in rhttpproxy entry in wcpsvc.yaml,
   # skip the patch.
   elif type(wcpsvc_yaml[_RHTTP_PROXY]) == int:
      logging.info('There is already rhttpproxy port value stored '
                   'in wcpsvc.yaml, skip the patch.')
      return

   # When there is rhttpproxy entry added in wcpsvc.yaml during wcpsvc.rpm
   # installation, store default port number in wcpsvc.yaml if it is unset.
   else:
      logging.info('Updating rhttpproxy port in wcpsvc config.')
      try:
         port = get_install_parameter(_RHTTP_PROXY_DEF_VAL,
                                      _RHTTPPROXY_PORT_DEFAULT)
         if not 1 <= int(port) <= 65535:
            raise Exception('Invalid port value {0}, must be between [1-65535]'
                            .format(port))
      except Exception as ex:
         logging.exception('Failed to retrieve rhttpproxy port {1}'.format(ex))
         raise

      try:
         fb = FileBuffer()
         fb.readFile(wcp_config_file_path)
         fb.findAndReplace("{" + _RHTTP_PROXY_DEF_VAL + "}", port)
         fb.writeFile(wcp_config_file_path)
      except Exception as ex:
         logging.exception('Failed to update rhttpproxy port in {0}: {1}'
                           .format(wcp_config_file_path, ex))
         raise


def doPatchingWithDependencies():
    pass


def doPatchingWithoutDependencies():
   """
   Do patching for wcpsvc.yaml file without dependent services. All patching
   functions for wcpsvc.yaml are called here.
   """
   preservePreviousEntries()
   substituteRhttpProxyPort()
   # Call your patching functions below


def doExpand(new_yaml, save_file=SAVED_WCPSVC_YAML):
   """
   Expand wcpsvc.yaml file
   """
   shutil.copyfile(WCPSVC_YAML, PRESERVED_WCPSVC_YAML)
   shutil.copyfile(new_yaml, save_file)
   preservePreviousEntries(PRESERVED_WCPSVC_YAML, save_file)
   substituteRhttpProxyPort(save_file)


def doRevert(save_file=SAVED_WCPSVC_YAML):
   """
   Revert changes from expand hook
   """
   if os.path.exists(save_file):
      os.remove(save_file)


def getChanges():
   return '%s: Changelog:\n%s' % (os.path.basename(__file__),
                                  '\n'.join(['* ' + line for line in
                                             CHANGELOG]))
