#! /usr/bin/env python3
#
#  known_hosts_update.py
#
# Copyright 2022-2023 VMware, Inc.  All rights reserved. -- VMware Confidential
# Description:  Provides class EdgeNodeVcfKnownHostsUtil.
# Disabled: True
#

import logging
import requests
import sys
import time
import urllib.parse
import urllib3

from vcf_utils.known_hosts_update import KnownHostsUpdate
from vcf_utils.vcf_creds import VcfCreds

urllib3.disable_warnings()


logger = logging.getLogger(__name__)


class EdgeNodeVcfKnownHostsUtil(object):
    """
    This class is used mainly to add an Edge node to VCF's known_hosts file.
    Since ssh access is required to obtain the node's key, this class also
    takes care of temporarily enabling ssh access to the Edge node if needed.
    If ssh is already enabled for the node, then it will be left enabled.
    """

    NSX_REL_URL_TNS = "api/v1/transport-nodes"
    TN_SSH_START = "start"
    TN_SSH_STOP = "stop"

    TN_SSH_STATE_RUNNING = "running"
    TN_SSH_STATE_STOPPED = "stopped"

    def __init__(self, nsxtCred, sddcMgrRest):
        """
        Setup to enable SDDC Manager access.

        nsxtCred : VcfCred
            Access info for NSX Manager hosting Edge node
        sddcMgrRest : SddcMgrRest
            Expected to be all set up & ready to use.
        edgeNodeNsxId : String
            NSX ID of Edge node whose ssh key we add to VCF known hosts
        """
        self.sddcMgrRest = sddcMgrRest

        self.nsxSession = requests.Session()
        self.nsxSession.auth = (nsxtCred.username,
                                nsxtCred.password)
        self.nsxSession.verify = False
        self.nsxMgrIp = nsxtCred.ipAddr
        self.nsxUrlPrefix = f"https://{self.nsxMgrIp}"

        return

    def _getTnSshStatus(self, edgeNodeNsxId):
        """
        Query
          api/v1/transport-nodes/.../node/services/ssh/status
        and report the runtime_state property of the returned
        NodeServiceStatusProperties instance.

        We expect runtime state to be either
        TN_SSH_STATE_RUNNING or TN_SSH_STATE_STOPPED.
        """

        relSvcStatusUrl = (
            "api/v1/transport-nodes/{}/node/services/ssh/status".format(
                edgeNodeNsxId))
        absSvcStatusUrl = urllib.parse.urljoin(self.nsxUrlPrefix,
                                               relSvcStatusUrl)
        resp = self.nsxSession.get(absSvcStatusUrl)

        resp.raise_for_status()

        reply = resp.json()
        replySchema = reply.get("_schema", "???")
        if replySchema != "NodeServiceStatusProperties":
            logger.error("! Unexpected ssh status reply schema: {}"
                         .format(replySchema))
            sys.exit(1)

        return reply.get("runtime_state", "")

    def _setTnSshState(self, edgeNodeNsxId, enable):
        """
        Return an absolute NSX URL to either start or stop the NSX service
        on the specified Edge transport node ID.

        enable : Boolean
            True to enable ssh, False to disable
        """

        if enable:
            action = "start"
            desired_state = "running"
        else:
            action = "stop"
            desired_state = "stopped"

        state = self._getTnSshStatus(edgeNodeNsxId)
        if state == desired_state:
            logger.info("Already in {} state".format(state))
            return True

        # form, following glcm's NsxtRawRestApiConnection.java, is this URL
        # plus a query parameter, action=start or action=stop

        relSvcSetUrl = "api/v1/transport-nodes/{}/node/services/ssh".format(
            edgeNodeNsxId)
        absSvcSetUrl = urllib.parse.urljoin(self.nsxUrlPrefix, relSvcSetUrl)

        absSvcStatusUrl = absSvcSetUrl + "/status"

        logger.info(f"* posting to url {absSvcSetUrl}")

        params = {"action": action}
        resp = self.nsxSession.post(absSvcSetUrl, params=params)
        if not resp:
            logger.error(f"! post failed: {resp}")
            resp.raise_for_status()

        # spin for a little while, watching to see if change takes effect:
        pollsRemaining = 5
        while pollsRemaining > 0:
            pollsRemaining -= 1

            time.sleep(2)
            logger.info(f"Get {absSvcStatusUrl}")
            state = self._getTnSshStatus(edgeNodeNsxId)

            logger.info("ssh runtime_state: {}".format(state))
            if state == desired_state:
                return True

        return False

    def freshenEdgeNodeInVcfKnownHosts(self, edgeNodeNsxId,
                                       hostUser, hostPass, dryrun=True):
        """
        Update specified Edge node's entry in VCF' known_hosts file.
        If necessary, we temporarily enable ssh to the Edge node
        in order to read its ssh key.
        """
        relUrl = (EdgeNodeVcfKnownHostsUtil.NSX_REL_URL_TNS + "/" +
                  edgeNodeNsxId)

        absUrl = urllib.parse.urljoin(self.nsxUrlPrefix, relUrl)

        # preload with initial state of Edge node before we touch it
        resp = self.nsxSession.get(absUrl)
        if not resp:
            logger.error("Failed to read Edge ID {} from NSX Manager: {}"
                         .format(edgeNodeNsxId, resp))
            resp.raise_for_status()
        edgeNodeTn = resp.json()

        ndi = edgeNodeTn["node_deployment_info"]
        settings = ndi["node_settings"]

        hostname = settings["hostname"]
        # orig_enable_ssh = settings["enable_ssh"]
        # orig_allow_ssh_root_login = settings["allow_ssh_root_login"]

        if dryrun:
            logger.info("Would freshen VCF known_hosts key for " + hostname)
        else:
            logger.info("Freshen VCF known_hosts key for " + hostname)

        curSshState = self._getTnSshStatus(edgeNodeNsxId)

        if (curSshState == EdgeNodeVcfKnownHostsUtil.TN_SSH_STATE_RUNNING):
            orig_enable_ssh = True
        else:
            orig_enable_ssh = False

        if not orig_enable_ssh:
            logger.warning("Temporarily enabling ssh to " + hostname)
            absUrl = self._setTnSshState(edgeNodeNsxId, True)

        # now do actual known_hosts update:
        khUpdater = KnownHostsUpdate(self.sddcMgrRest)

        freshenKeyResult = khUpdater.freshenHostKey(hostname,
                                                    hostUser, hostPass,
                                                    dryrun=dryrun)

        # before checking freshen status, restore ssh disable if necessary
        if not orig_enable_ssh:
            logger.warning("Re-disabling ssh to " + hostname)
            absUrl = self._setTnSshState(edgeNodeNsxId, False)

        if not freshenKeyResult:
            logger.error("Failed to freshen known_hosts entry for {}"
                         .format(hostname))
            sys.exit(1)

        return


if __name__ == '__main__':
    # (by default, logger is only enabled for warning and higher)
    logging.basicConfig(level=logging.INFO, format='%(message)s')

    import argparse
    import os

    default_user = os.getenv("SDDC_SSO_USERNAME")
    default_pass = os.getenv("SDDC_SSO_PASSWORD")

    parser = argparse.ArgumentParser(
        description="Test VCF known_hosts util")
    parser.add_argument(
        '--workload', '-w', required=False,
        help='name of VCF workload domain to use (case sensitive)')
    parser.add_argument(
        '--user', '-u', required=False, default=default_user,
        help='Name of SDDC SSO admin user to authenticate as')
    parser.add_argument(
        '--password', '-p', required=False, default=default_pass,
        help='Password for specified user')
    args = parser.parse_args()

    creds = VcfCreds(workload=args.workload,
                     username=args.user, password=args.password)

    from vcf_utils.sddc_mgr_rest import SddcMgrRest

    # run either on SDDC Manager, or with password-free ssh access to it:
    sddcMgrRest = SddcMgrRest()

    EDGE_NODE_ID = "4b2bc2ce-eba7-40ef-98bf-f92a587f8ef8"
    enUser = "root"
    enPass = "xxx"

    freshener = EdgeNodeVcfKnownHostsUtil(creds.nsxtCred, sddcMgrRest)

    freshener.freshenEdgeNodeInVcfKnownHosts(EDGE_NODE_ID, enUser, enPass,
                                             dryrun=False)
