#! /usr/bin/env python3
#
#  vcf_creds.py
#
# Copyright 2020 VMware, Inc.  All rights reserved. -- VMware Confidential
# Description:  Provides class VcfCreds
# Disabled: True
#

import argparse
import distutils.util
import getpass
import logging
import os
import pprint
import sys
import time
import traceback

from datetime import datetime, timedelta, timezone
from vcf_utils.sddc_mgr_rest import SddcMgrRest
from vcf_utils.vcf_cred import VcfCred


logger = logging.getLogger(__name__)


class VcfCreds:
    """
    Retrieves credentials info from SDDC Manager.
    At present, we obtain only NSX-T Manager and vCenter access info.
    But this could likely be expanded fairly easily.
    """

    # Maximum minutes since last auth token refresh (or creation), before
    # we must refresh the token. SDDC Manager's token API gives us an hour,
    # but we knock that down to 55 minutes for a safety margin.
    MAX_MINUTES_BEFORE_TOKEN_REFRESH = 55

    # entity types which might be used when submitting a credential in order
    # to obtain an auth token.
    CRED_ENTITY_TYPES = [
        "ESXI", "VCENTER", "PSC", "NSX_MANAGER", "NSX_CONTROLLER", "NSX_EDGE",
        "NSXT_MANAGER", "NSXT_EDGE", "VRLI", "VROPS", "VRA", "VRSLCM",
        "BACKUP", "VXRAIL_MANAGER", "AD"]

    def _checkIfTokenSecurityEnabled(self):
        """Return a simple True or False value according to whether
        'feature.vcf.public.api.security' is enabled or not.
        """
        API_SEC_FEATURE_NAME = 'feature.vcf.public.api.security'

        #  Default to on, since API security has now been onboarded
        #  so removal of the feature-toggle implies always-on.
        apiSecurityEnabled = True

        try:
            featureEnables = self.sddcMgrRest.runSddcMgrRestGet(
                "http://127.0.0.1:7100/features/list")

            if API_SEC_FEATURE_NAME in featureEnables:
                apiSecurityEnabled = (
                    bool(distutils.util.strtobool(
                        featureEnables[API_SEC_FEATURE_NAME])))

        except Exception as e:
            logger.error("Failed to check API security enable: %s" % e)
            sys.exit(1)

        return apiSecurityEnabled

    def _createAuthToken(self, username, password, sddcMgrIp="localhost"):
        """
        Given a username and password, do the initial "bearer token"
        retrieval for this user to access VCF public APIs.
        This method should only be called once over the lifetime of
        a given script run.

        We save the token and its refresh info in our instance data.
        If code wants to actually use the token, it should call
        _getAuthToken() to obtain it. This ensures that the returned token
        has been first refreshed if it was close to expiration.

        This gives us a total working time for a single _createAuthToken()-
        created auth token (& its refreshes) of 24 hours. Which presumably
        suffices for present uses.
        """

        url = 'https://{}/v1/tokens'.format(sddcMgrIp)
        try:
            logger.info('Getting Bearer token')
            logger.debug('URL : POST : {}'.format(url))
            payload = ('{"username": "' + username +
                       '", "password": "' + password + '"}')
            tokenReply = self.sddcMgrRest.runSddcMgrRestPostString(
                url, payloadString=payload)

            self.__saveAuthToken(tokenReply['accessToken'])
            self.refreshToken = tokenReply['refreshToken']['id']
            self.tokenApiAddr = sddcMgrIp

        except Exception as ex:
            logger.error('Failed to get bearer token!')
            pp = pprint.PrettyPrinter()
            logger.error("Response:\n%s" % pp.pformat(tokenReply))
            logger.error(ex)
            logger.error(traceback.format_exc())
            sys.exit(1)

        return

    def __refreshAuthToken(self):
        """
        Refresh our current VCF authentication (access) token. This relies on
        the refresh token already saved by _createAuthToken().
        """

        url = 'https://{}/v1/tokens/access-token/refresh'.format(
            self.tokenApiAddr)
        try:
            logger.info('Refreshing Bearer token')
            logger.debug('URL : PATCH : {}'.format(url))
            payload = '"{}"'.format(self.refreshToken)
            refreshReply = self.sddcMgrRest.runSddcMgrRestPatchString(
                url, payloadString=payload)

            self.__saveAuthToken(refreshReply)

        except Exception as ex:
            logger.error('Failed to refresh bearer token!')
            pp = pprint.PrettyPrinter()
            logger.error("Response:\n%s" % pp.pformat(refreshReply))
            logger.error(ex)
            logger.error(traceback.format_exc())
            sys.exit(1)

        return

    def __saveAuthToken(self, authToken):
        """
        Common way of preserving & timestamping an auth token.
        Use __getAuthToken() to retrieve the token, it will auto-refresh
        if needed.
        """
        self.authToken = authToken
        self.lastTokenRefreshTime = datetime.now(timezone.utc)

        return

    def _getAuthToken(self):
        """
        Retrieve a usable VCF auth token. This is based on the one first
        saved by _createAuthToken(), but if necessary it is refreshed to
        assure that the token we return is still usable.

        This gives a useful working window of 24 hours from the time that
        _createAuthToken() was called.
        """
        elapsedSinceRefresh = (datetime.now(timezone.utc) -
                               self.lastTokenRefreshTime)

        minsSinceRefresh = round(elapsedSinceRefresh / timedelta(minutes=1))
        if minsSinceRefresh > VcfCreds.MAX_MINUTES_BEFORE_TOKEN_REFRESH:
            self.__refreshAuthToken()

        return self.authToken

    def _getCredsWithTokenAuth(self):

        # Get domain inventory, which is where entity IDs are found
        domId, domInv = self._getDomInventory()

        # save VCF inventory ID of WLD which we're operating on
        self.vcfDomainId = domId

        wldVcenterInfo = self._getWldVcenterInfo(domId)

        # grab count of workload domains supported by this NSX-T cluster
        self.wldsSupportedByOurNsxt = len(
            domInv["nsxtClusters"][0]["domainIds"])
        logger.info("count of WLDs supported by our NSX-T cluster: %s" %
                    self.wldsSupportedByOurNsxt)

        # find NSX-T cred for our NSX-T cluster's ID:
        self.nsxtClusterId = domInv["nsxtClusters"][0]["id"]

        pscUrl = "http://localhost/v1/credentials?resourceType=PSC"
        pscCredsInfo = self.sddcMgrRest.runSddcMgrRestGet(
            pscUrl, authToken=self._getAuthToken())

        # we expect only a single PSC / SSO record:
        pscCredCount = len(pscCredsInfo["elements"])

        if pscCredCount == 1:
            rawPscCred = pscCredsInfo["elements"][0]
        else:
            rawPscCred = self._getCurDomPscCred(domId, wldVcenterInfo,
                                                pscCredsInfo)

        # we need to find the correct VC address for the target WLD:
        vcIp = wldVcenterInfo["managementIpAddress"]

        self.vcCred = VcfCred(username=rawPscCred["username"],
                              password=rawPscCred["password"],
                              ipAddr=vcIp)

        nsxtUrl = "http://localhost/v1/credentials?resourceType=NSXT_MANAGER"
        nsxtCredsInfo = self.sddcMgrRest.runSddcMgrRestGet(
            nsxtUrl, authToken=self._getAuthToken())

        # typically multiple NSX-T entries, find match for our workload
        for newCred in nsxtCredsInfo["elements"]:
            if newCred["credentialType"] == "API":
                resInfo = newCred["resource"]
                if resInfo["resourceId"] == self.nsxtClusterId:
                    self.nsxtCred = VcfCred(username=newCred["username"],
                                            password=newCred["password"],
                                            ipAddr=resInfo["resourceIp"])
                    break

        if not self.nsxtCred:
            logger.error("Failed to find NSX-T cred matching %s" %
                         self.nsxtClusterId)

        return

    def _getCurDomPscCred(self, domId, wldVcenterInfo, pscCredsInfo):
        """
        Given a VCF workload domain ID, the workload's vCenter info, and
        a list of available PSC credentials, find the correct cred.

        Per the Isolated Workload Domains Functional Spec, a workload
        either has its own PSC or else it uses the MGMT PSC. So this
        isn't (currently at least) as complex as one might imagine.
        """
        pscCreds = pscCredsInfo.get("elements")
        if not pscCreds:
            return None

        if len(pscCreds) == 1:
            return pscCreds[0]

        targetPscId = wldVcenterInfo["pscId"]

        for cred in pscCreds:
            pscId = cred["resource"]["resourceId"]
            if pscId == targetPscId:
                logger.info("Using workload-specific PSC ID " + pscId)
                return cred

        # no specific match, so default to the MANAGEMENT PSC

        # first, find the MGMT PSC
        mgmtPscId = None
        for vc in self.vcList:
            if vc["domainType"] == "MANAGEMENT":
                mgmtPscId = vc["pscId"]

        mgmtCred = None
        for cred in pscCreds:
            pscId = cred["resource"]["resourceId"]
            if pscId == mgmtPscId:
                mgmtCred = cred
                break

        return mgmtCred

    def getCredsByIpWithTokenAuth(self, mgmtIp, credType):
        """
        For a given CRED_ENTITY_TYPES resource type value, look up and return
        all credentials of that type, which match the specified mgmtIp.

        The user may wish to refresh our auth token first, if this is part of
        a long-running operation.

        mgmtIp : String
            IP address of entity to retrieve credentials for.
        credType : String
            Resource type of entity to retrieve credentials for.
            Typically one of the values from CRED_ENTITY_TYPES.
        """

        credsUrl = "http://localhost/v1/credentials?resourceType=" + credType
        credsInfo = self.sddcMgrRest.runSddcMgrRestGet(
            credsUrl, authToken=self._getAuthToken())

        matchingCreds = []

        # typically multiple NSX-T entries, find match for our workload
        for newCred in credsInfo["elements"]:
            resIp = newCred["resource"]["resourceIp"]
            if resIp != mgmtIp:
                continue

            matchingCred = VcfCred(username=newCred["username"],
                                   password=newCred["password"],
                                   ipAddr=resIp,
                                   credType=newCred["credentialType"])
            matchingCreds.append(matchingCred)

        if not matchingCreds:
            logger.error(f"Failed to find {credType} cred matching {mgmtIp}")

        return matchingCreds

    def _getWldVcenterInfo(self, domId):
        """Get VCF inventory entry for vCenter used with the supplied
        VCF WLD ID.
        """
        vcList = self.sddcMgrRest.runSddcMgrRestGet(
            "http://127.0.0.1/inventory/vcenters")

        """Returns a list of entries like so:
        [
           {
              "bundleRepoDatastore" : "lcm-bundle-repo",
              "datastoreForVmDeploymentName" : "sfo01-m01-vsan",
              "hostName" : "vi-vcenter-1.vrack.vsphere.local",
              "vmName" : "vi-vcenter-1",
              "status" : "ACTIVE",
              "domainId" : "5623e3e2-8185-40ec-bee8-2291d040a735",
              "id" : "88205cbc-0b16-4bc4-b809-daee62bbf0a0",
              "version" : "7.0.0.10000-15743180",
              "domainType" : "VI",
              "managementIpAddress" : "10.0.0.20"
           }
        ]
        """
        # save for possible later use matching [isolated] workloads
        self.vcList = vcList

        for vc in vcList:
            if domId == vc["domainId"]:
                return vc

        return None

    def _getDomList(self):
        """Get list of domains in SDDC Manager's inventory."""
        domList = self.sddcMgrRest.runSddcMgrRestGet(
            "http://127.0.0.1/inventory/domains")

        """Returns a list of entries like so:
        [
           {
              "status" : "ACTIVE",
              "id" : "3e1a866c-c7db-418c-a85a-2b97eaecc2fe",
              "name" : "MGMT",
              "organization" : "COM",
              "type" : "MANAGEMENT"
           }
        ]
        """
        return domList

    def _getDomInventory(self):
        """Grab full inventory for a single domain.  If there's only one,
        or if the user specifies a target domain, we return that domain's
        inventory. Otherwise we fail.

        self.workload_name - Name of domain to return. If not specified,
            we fail unless there is only a single domain in the system.
            If this is None, we set it to the inferred default value.

        We return a tuple of domainId, domainInventory.
        """

        domList = self._getDomList()

        domainCount = len(domList)
        if domainCount == 0:
            logger.error("No domain info found in SDDC Manager's inventory")
            sys.exit(1)

        domIndex = 0  # until found otherwise

        if domainCount > 1:
            domainNames = []
            for dom in domList:
                domainNames.append(dom["name"])

            if self.workload_name and self.workload_name in domainNames:
                domIndex = domainNames.index(self.workload_name)
            elif self.workload_name:
                logger.error("Unknown workload name %s, please specify "
                             "one of" % self.workload_name)
                logger.error("  %s" % domainNames)
                sys.exit(1)
            else:
                logger.error("Found %s workloads: %s" %
                             (domainCount, domainNames))
                logger.error("Please use the --workload option to select "
                             "one of them")
                sys.exit(1)
        else:
            # domIndex will likely always be zero here, but use it anyway
            domNameInInv = domList[domIndex]["name"]
            if not self.workload_name:
                self.workload_name = domNameInInv
            elif self.workload_name != domNameInInv:
                logger.error("Unknown workload name %s, only known "
                             "workload is %s" %
                             (self.workload_name, domNameInInv))
                sys.exit(1)

        domId = domList[domIndex]["id"]

        # grab inventory for workload domain
        domInv = self.sddcMgrRest.runSddcMgrRestGet(
            "http://127.0.0.1/inventory/domains/%s/inventory" % domId)

        return domId, domInv

    def __init__(self, username, password, workload=None,
                 sddcMgrAddr="10.0.0.4"):
        """Retrieve a snapshot of credentials, which we hold for the life
        of our instance.

        If present, workload is the name of the VCF WLD to use.

        If present, sddcMgrAddr is the address of the external SDDC Manager
        VM which we assume password-free ssh access to as user "vcf".
        Ignored if we are running inside an SDDC Manager VM."""

        if not username:
            logger.error("Username is required.")
            sys.exit(1)

        if not password:
            logger.warning("Please enter password for user %s" % username)
            password = getpass.getpass()

        self.workload_name = workload

        self.sddcMgrRest = SddcMgrRest(sddcMgrAddr=sddcMgrAddr)

        self.tokenSecurityEnabled = self._checkIfTokenSecurityEnabled()
        if self.tokenSecurityEnabled:
            self._createAuthToken(username, password)
            self._getCredsWithTokenAuth()
        else:
            logger.error("Only systems with token authentication are "
                         "supported.")
            sys.exit(1)

        logger.info("workload_name = %s" % self.workload_name)


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

    logger.info("VcfCreds test running..")

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

    parser = argparse.ArgumentParser(
        description="Test VCF inventory / credentials retrieval")
    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)

    pp = pprint.PrettyPrinter()

    logger.info("nsxt cred:\n {0}".format(creds.nsxtCred))
    logger.info("vc cred:\n {0}".format(creds.vcCred))
    logger.info("WLD id: {0}\n".format(creds.vcfDomainId))

    # now force a token refresh by setting the last-refresh timestamp
    # to be older that our refresh threshold
    oldTimestamp = (time.time() -
                    60 * (VcfCreds.MAX_MINUTES_BEFORE_TOKEN_REFRESH + 3))
    creds.lastTokenRefreshTime = datetime.fromtimestamp(oldTimestamp,
                                                        timezone.utc)

    nsxtCreds = creds.getCredsByIpWithTokenAuth("10.0.0.30", "NSXT_MANAGER")
    logger.info("MGMT NSX-T creds:")
    for cred in nsxtCreds:
        logger.info("  {}".format(cred))
