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

import json
import logging


logger = logging.getLogger(__name__)


class KnownHostsUpdate(object):
    """
    *Internal class* In general, please use EdgeNodeVcfKnownHostsUtil instead.

    Utility used for adding or removing host entries in the VCF known_hosts
    file. When doing an add, this utility assumes that ssh access to the
    target host is already enabled.
    """

    APPLIANCE_KNOWN_HOSTS_URL = ("http://localhost/appliancemanager/ssh/"
                                 "knownHosts")

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

        sddcMgrRest : SddcMgrRest
            Expected to be all set up & ready to use.
        """
        self.sddcMgr = sddcMgrRest
        return

    def readCurrentKH(self):

        url = KnownHostsUpdate.APPLIANCE_KNOWN_HOSTS_URL
        curKnownHosts = self.sddcMgr.runSddcMgrRestGet(url)
        return curKnownHosts

    def removeHost(self, curKnownHosts, hostFqdn):
        """
        Given a known_hosts spec returned by readCurrentKH(),
        we return a copy of the spec but with all keys for
        hostFqdn removed.
        """
        hostFqdnCf = hostFqdn.casefold()
        cleanedKeyList = []
        for keySpec in curKnownHosts.get("knownHosts", []):
            host = keySpec.get("host", "???")
            keyType = keySpec.get("keyType", "???")
            if hostFqdnCf == host.casefold():
                logger.info("dropping old {}, key type {}".format(host,
                                                                  keyType))
            else:
                cleanedKeyList.append(keySpec)

        return {"knownHosts": cleanedKeyList}

    def _tryLocalKeyscan(self, hostFqdn):
        try:
            scanCmd = ['ssh-keyscan', '-4', '-t', 'rsa', hostFqdn]
            scanResultsByte = self.sddcMgr._runSddcMgrCmd(scanCmd)
            # convert to UTF-8 and remove any trailing newline
            scanResults = scanResultsByte.decode('utf-8').strip()
        except Exception as e:
            # be a little more friendly than basic subprocess.check_output()
            scanResults = str(e)

        return scanResults

    def _tryRemoteKeyscan(self, hostFqdn, hostUser, hostPassword):
        """
        As a fallback, try ssh'ing into the target Edge node and running
        its local ssh-keyscan utility instead of our own.
        """

        # scan command to run inside Edge node:
        rmtScanCmd = 'ssh-keyscan -4 -t rsa localhost'

        client = None
        _stdin = None

        try:
            # We only import paramiko on demand, so we don't rely on its
            # presence unless our local ssh-keyscan is inadequate.
            import paramiko

            client = paramiko.client.SSHClient()
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            client.connect(hostFqdn, username=hostUser, password=hostPassword)
            _stdin, _stdout, _stderr = client.exec_command(rmtScanCmd)
            scanResults = _stdout.read().decode()

        except Exception as e:
            scanResults = str(e)

        finally:
            # https://docs.paramiko.org/en/stable/api/client.html
            # strongly advises an explicit close:
            if client is not None:
                client.close()

                # and close() alone can still intermittently fall prey to
                # https://github.com/paramiko/paramiko/issues/1078
                # so we also follow
                # https://github.com/paramiko/paramiko/issues/1078#issuecomment-596771584
                del client
                if _stdin is not None:
                    del _stdin, _stdout, _stderr

        return scanResults

    def _decodeKeyscanResults(self, scanResults, hostFqdn):
        """
        Given raw text output from an ssh-keyscan call, extract fields
        of interest. Returns each field as None when decode fails.

        Return: keyFqdn, keyType, keyValue
        """
        keyParts = scanResults.split()
        if ((len(keyParts) < 3) or
                (keyParts[0].casefold() != hostFqdn.casefold()) or
                (keyParts[1] != 'ssh-rsa')):
            keyFqdn = None
            keyType = None
            keyValue = None
        else:
            keyFqdn = keyParts[0]
            keyType = keyParts[1]
            keyValue = keyParts[2]

        return keyFqdn, keyType, keyValue

    def getCurHostKey(self, hostFqdn, hostUser, hostPassword):
        """
        Given an FQDN, scan it and return the resulting key info as a
        "key spec" per VCF's known_hosts spec format.

        Assumes that ssh access to the target host has already been enabled.

        Right now, we read only RSA keys since that is what NSX Edge nodes
        use, and what
            vcf-common-core/vcf-secure-libs/vcf-secure-ssh-lib/src/main/
                java/com/vmware/evo/sddc/common/util/SshKeyHelper.java
        limits itself to retrieving.
        """

        # NB: might also use paramiko here, but not sure what range
        #     of SDDC Manager versions offer it.

        scanResults = self._tryLocalKeyscan(hostFqdn)

        keyFqdn, keyType, keyValue = self._decodeKeyscanResults(scanResults,
                                                                hostFqdn)

        if not keyType:
            logger.warning("Failed to read RSA key for {} with local "
                           "ssh-keyscan: {}".format(hostFqdn, scanResults))

            scanResults = self._tryRemoteKeyscan(hostFqdn,
                                                 hostUser, hostPassword)

            # since _tryRemoteKeyscan() ran against Edge node's own localhost
            # we need to adjust our expected keyFqdn handling.
            keyFqdn, keyType, keyValue = self._decodeKeyscanResults(
                scanResults, "localhost")
            # but we still want to report the Edge node's true fqdn
            keyFqdn = hostFqdn

            if not keyType:
                logger.warning("Failed to read RSA key for {} with remote "
                               "ssh-keyscan: {}".format(hostFqdn,
                                                        scanResults))
                return None

        return {"host": keyFqdn,
                "keyType": keyType,
                "key": keyValue}

    def freshenHostKey(self, hostFqdn, hostUser, hostPassword,
                       dryrun=True):
        """
        Read VCF appliance's known_hosts file, remove all current key
        entries for the specified hostFqdn, and replace them with a
        current key read directly from the host.
        """

        logger.info("Freshen known_hosts key for " + hostFqdn)

        oldKnownHosts = self.readCurrentKH()

        curHostKey = self.getCurHostKey(hostFqdn, hostUser, hostPassword)
        if not curHostKey:
            logger.error("Failed to freshen known_hosts for {}"
                         .format(hostFqdn))
            return False

        logger.info("Current ssh key obtained for " + hostFqdn)

        if dryrun:
            logger.info("In dry-run, so not updating known_hosts for "
                        + hostFqdn)
            return True

        newKnownHosts = self.removeHost(oldKnownHosts, hostFqdn)

        newKnownHosts["knownHosts"].append(curHostKey)

        url = KnownHostsUpdate.APPLIANCE_KNOWN_HOSTS_URL
        ret = self.sddcMgr.runSddcMgrRestPostString(url, newKnownHosts)

        logger.info("Ran post, result = {}".format(ret))

        return True


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

    import sys
    sys.path.append("../pyVpx")
    sys.path.append("../third-party")
    sys.path.append("..")

    from vcf_utils.sddc_mgr_rest import SddcMgrRest

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

    # art-edge-cluster (lanier-goat/conf/sp/wcp/art-wcp-edge.json)
    fqdnToFreshen = "nsxt-edge-node-12.vrack.vsphere.local"
    hostUser = "ssher"
    hostPassword = "xxx"

    if not khUpdater.freshenHostKey(fqdnToFreshen, hostUser, hostPassword):
        print("! failed to freshen known_hosts for " + fqdnToFreshen)
        sys.exit(1)

    # we think it worked, do a quick sanity-check
    curKnownHosts = khUpdater.readCurrentKH()
    curKey = khUpdater.getCurHostKey(fqdnToFreshen, hostUser, hostPassword)

    khKey = None
    for keySpec in curKnownHosts["knownHosts"]:
        if keySpec["host"].casefold() != fqdnToFreshen.casefold():
            continue

        if khKey:
            print("! found multiple entries in known_hosts for {}"
                  .format(fqdnToFreshen))
        else:
            khKey = keySpec

    if curKey == khKey:
        print("Key in known hosts is correct for " + fqdnToFreshen)
    else:
        print("! Key in known_hosts doesn't match for " + fqdnToFreshen)
        print("  From known_hosts: " + json.dumps(khKey, indent=3))
        print("  From live host: " + json.dumps(curKey, indent=3))
