hcs-core 0.1.283__py3-none-any.whl → 0.1.316__py3-none-any.whl

hcs_core/sglib/auth.py

@@ -22,7 +22,7 @@ import time
 import jwt
 from authlib.integrations.httpx_client import OAuth2Client
 
-from hcs_core.ctxp import CtxpException, panic, profile
+from hcs_core.ctxp import CtxpException, profile
 from hcs_core.ctxp.jsondot import dotdict, dotify
 
 from .csp import CspClient
@@ -30,122 +30,141 @@ from .csp import CspClient
 log = logging.getLogger(__name__)
 
 
-def _get_profile_auth_hash(effective_profile):
-    csp = effective_profile.csp
-    text = json.dumps(csp, default=vars)
-    return profile.name() + "#" + hashlib.md5(text.encode("ascii"), usedforsecurity=False).hexdigest()
-
-
-def _is_auth_valid(auth_data, effective_profile):
-    if not auth_data:
-        return
-    if not auth_data.token:
-        return
-    if auth_data.hash != _get_profile_auth_hash(effective_profile):
-        return
+def _is_auth_valid(auth_data):
     leeway = 60
-    if time.time() + leeway >= auth_data.token.expires_at:
-        return
-    return True
-
-
-def _decode_http_basic_auth_token(basic_token: str):
-    import base64
-
-    try:
-        decoded = base64.b64decode(basic_token).decode("utf-8")
-        client_id, client_secret = decoded.split(":")
-        return client_id, client_secret
-    except Exception as e:
-        raise CtxpException(f"Invalid basic http auth token: {e}")
+    return auth_data and time.time() + leeway < auth_data["expires_at"]
 
 
 _login_lock = threading.Lock()
 
 
-def login(force_refresh: bool = False, panic_on_failure: bool = True):
+def login(force_refresh: bool = False, verbose: bool = False):
     """Ensure login state, using credentials from the current profile. Return oauth token."""
+    return _populate_token_with_cache(profile.current().csp, force_refresh, verbose)
 
-    _login_lock.acquire()
-    try:
-        effective_profile = profile.current()
-
-        auth_data = profile.auth.get()
-        if force_refresh or not _is_auth_valid(auth_data, effective_profile):
-            oauth_token = _get_new_oauth_token(auth_data.token, effective_profile)
-            if oauth_token:
-                use_oauth_token(oauth_token, effective_profile)
-            elif panic_on_failure:
-                panic(
-                    "Login failed. If this is configured API key or client credential, refresh the credential from CSP and update profile config. If this is browser based interactive login, login again."
-                )
-            else:
-                return None
-        else:
-            oauth_token = auth_data.token
-        return oauth_token
-    finally:
-        _login_lock.release()
-
-
-def _get_new_oauth_token(old_oauth_token, effective_profile):
-    csp_config = effective_profile.csp
-
-    csp_client = CspClient(url=csp_config.url)
-
-    if csp_config.apiToken:
-        oauth_token = csp_client.login_with_api_token(csp_config.apiToken)
-    elif csp_config.clientId:
-        oauth_token = csp_client.login_with_client_id_and_secret(
-            csp_config.clientId, csp_config.clientSecret, csp_config.orgId
-        )
-    elif csp_config.basic:
-        client_id, client_secret = _decode_http_basic_auth_token(csp_config.basic)
-        oauth_token = csp_client.login_with_client_id_and_secret(client_id, client_secret, csp_config.orgId)
-    else:
-        # This should be a config from interactive login.
-        # Use existing oauth_token to refresh.
-        if not old_oauth_token:
-            old_oauth_token = profile.auth.get().token
-        if old_oauth_token:
-            try:
-                oauth_token = _refresh_oauth_token(old_oauth_token, csp_config.url)
-            except Exception as e:
-                oauth_token = None
-                log.warning(e)
-        else:
-            oauth_token = None
-
-    if oauth_token and not oauth_token.get("expires_at"):
-        oauth_token["expires_at"] = int(time.time() + oauth_token["expires_in"])
-
-    return oauth_token
-
-
-def _refresh_oauth_token(old_oauth_token: dict, csp_url: str):
+
+def refresh_oauth_token(old_oauth_token: dict, csp_url: str):
     with OAuth2Client(token=old_oauth_token) as client:
         log.debug("Refresh auth token...")
-        new_token = client.refresh_token(csp_url + "/csp/gateway/am/api/auth/token")
+        token_url = csp_url + "/csp/gateway/am/api/auth/token"
+        from .login_support import identify_client_id
+
+        csp_specific_req_not_oauth_standard = (identify_client_id(csp_url), "")
+        new_token = client.refresh_token(token_url, auth=csp_specific_req_not_oauth_standard)
         log.debug(f"New auth token: {new_token}")
         if not new_token:
             raise Exception("CSP auth refresh failed.")
         if "cspErrorCode" in new_token:
             raise Exception(f"CSP auth failed: {new_token.get('message')}")
+
     return new_token
 
 
-class MyOAuth2Client(OAuth2Client):
-    def __init__(self):
+def _has_credential(auth_config: dict):
+    return (
+        auth_config.get("apiToken")
+        or auth_config.get("api_token")
+        or auth_config.get("clientId")
+        or auth_config.get("client_id")
+        or auth_config.get("basic")
+    )
+
+
+def get_auth_cache(auth_config: dict):
+    cache, hash, token = _get_auth_cache(auth_config)
+    return token
+
+
+def _get_auth_cache(auth_config: dict):
+    text = json.dumps(auth_config, default=vars)
+    hash = hashlib.md5(text.encode("ascii"), usedforsecurity=False).hexdigest()
+    auth_cache = profile.auth.get()
+    token = auth_cache.get(hash, None)
+    if token and not _is_auth_valid(token):
+        token = None
+    return auth_cache, hash, token
+
+
+def save_auth_cache(auth_config: dict, token: dict):
+    """Save the auth token to the profile auth cache."""
+    cache, hash, _ = _get_auth_cache(auth_config)
+    if not token:
+        if hash in cache:
+            del cache[hash]
+        return
+    if not token.get("expires_at"):
+        token["expires_at"] = int(time.time() + token["expires_in"])
+    cache[hash] = token
+    profile.auth.set(cache)
+
+
+def _populate_token_with_cache(auth_config: dict, force_refresh: bool = False, verbose: bool = False):
+    if verbose:
+        print("_populate_token_with_cache")
+    with _login_lock:
+        cache, hash, token = _get_auth_cache(auth_config)
+        if token and not force_refresh:
+            if verbose:
+                print("Using cached auth token.")
+            return token
+
+        # invalid token. Refresh or recreate it.
+        if token:
+            # try using refresh token if possible
+            if auth_config.get("provider", "vmwarecsp") == "vmwarecsp":
+                if verbose:
+                    print("Provider: vmwarecsp.")
+                try:
+                    token = refresh_oauth_token(token, auth_config.url)
+                except Exception as e:
+                    log.debug(f"Failed to refresh OAuth token: {e}")
+                    token = None
+            else:
+                # hcs auth-service. Does not support refresh token.
+                if verbose:
+                    print("Provider: hcs auth-service.")
+                token = None
+
+        if not token:
+            if _has_credential(auth_config):
+                if verbose:
+                    print("Config:", auth_config)
+                token = CspClient.create(**auth_config).oauth_token()
+                if verbose:
+                    print("Token:", json.dumps(token, indent=4))
+            else:
+                if auth_config.get("browser"):
+                    from .login_support import login_via_browser
+
+                    token = login_via_browser(auth_config.url, auth_config.orgId)
+                    if not token:
+                        raise CtxpException("Browser auth failed.")
+                else:
+                    raise CtxpException("Browser auth was never attempted and no client credentials or API token provided.")
+
+        if not token.get("expires_at"):
+            token["expires_at"] = int(time.time() + token["expires_in"])
+
+        cache[hash] = token
+        profile.auth.set(cache)
+    return token
+
+
+class CustomOAuth2Client(OAuth2Client):
+    def __init__(self, auth_config: dict):
         super().__init__()
+        self.auth_config = auth_config
 
     def ensure_token(self):
         # pylint: disable=access-member-before-definition
         if self.token is None or not super().ensure_active_token():
-            self.token = login()
+            self.token = _populate_token_with_cache(self.auth_config)
 
 
-def oauth_client():
-    return MyOAuth2Client()
+def oauth_client(auth_config: dict = None):
+    if not auth_config:
+        auth_config = profile.current().csp
+    return CustomOAuth2Client(auth_config)
 
 
 def details(get_org_details: bool = False) -> dotdict:
@@ -174,9 +193,3 @@ def details_from_token(oauth_token, get_org_details: bool = False):
 def get_org_id_from_token(oauth_token: str) -> str:
     decoded = jwt.decode(oauth_token["access_token"], options={"verify_signature": False})
     return decoded["context_name"]
-
-
-def use_oauth_token(oauth_token, effective_profile=None):
-    if effective_profile is None:
-        effective_profile = profile.current()
-    profile.auth.set({"token": oauth_token, "hash": _get_profile_auth_hash(effective_profile)})

hcs_core/sglib/cli_options.py

@@ -18,7 +18,21 @@ import os
 import click
 
 from hcs_core.ctxp import CtxpException, recent
-from hcs_core.ctxp.cli_options import *
+from hcs_core.ctxp.cli_options import apply_env as apply_env
+from hcs_core.ctxp.cli_options import confirm as confirm
+from hcs_core.ctxp.cli_options import env as env
+from hcs_core.ctxp.cli_options import exclude_field as exclude_field
+from hcs_core.ctxp.cli_options import field as field
+from hcs_core.ctxp.cli_options import first as first
+from hcs_core.ctxp.cli_options import force as force
+from hcs_core.ctxp.cli_options import formatter as formatter
+from hcs_core.ctxp.cli_options import ids as ids
+from hcs_core.ctxp.cli_options import limit as limit
+from hcs_core.ctxp.cli_options import output as output
+from hcs_core.ctxp.cli_options import search as search
+from hcs_core.ctxp.cli_options import sort as sort
+from hcs_core.ctxp.cli_options import verbose as verbose
+from hcs_core.ctxp.cli_options import wait as wait
 
 org_id = click.option(
     "--org",

hcs_core/sglib/client_util.py

@@ -14,95 +14,178 @@ limitations under the License.
 """
 
 import json
-import os
+import logging
 import sys
 import threading
 import time
 from typing import Callable, Iterator
 
 from hcs_core.ctxp import CtxpException, panic, profile
-from hcs_core.sglib import hcs_client
 from hcs_core.sglib.ez_client import EzClient
 from hcs_core.util import duration, exit
 from hcs_core.util.query_util import PageRequest, with_query
 
+log = logging.getLogger(__name__)
+
 _caches = {}
 
 _client_instance_lock = threading.RLock()
 
 
-def hdc_service_client(service_name: str) -> EzClient:
-    _client_instance_lock.acquire()
-    try:
-        instance = _caches.get(service_name)
-        if not instance:
+def _get_service_override(service_name: str):
+    profile_data = profile.current()
+    override = profile_data.get("override", {})
+    service_override = {}
+    for k, v in override.items():
+        if service_name == k.lower():
+            service_override = v
 
-            def _get_url():  # make it deferred so no need to initialize profile
-                url = _get_hcs_url_considering_env_override(service_name)
-                if not url.endswith("/"):
-                    url += "/"
-                url += service_name
-                return url
+    if not service_override:
+        if service_name == "org-service":
+            service_override = override.get("org", {})
+        elif service_name.find("-") >= 0:
+            camel_name = "".join(word.capitalize() for word in service_name.split("-"))
+            camel_name = camel_name[0].lower() + camel_name[1:]
+            service_override = override.get(camel_name, {})
+    return service_override
 
-            instance = hcs_client(_get_url)
-            _caches[service_name] = instance
-        return instance
-    finally:
-        _client_instance_lock.release()
-
-
-def _get_hcs_url_considering_env_override(service_name: str):
-    # service_name = service_name.replace("-", "_")
-    # env_name = f"hcs_{service_name}_url"
-    # url = os.environ.get(env_name)
-    # if url:
-    #     print(f"Using env override for {env_name}: {url}")
-    #     return url
-    # env_name = env_name.upper()
-    # url = os.environ.get(env_name)
-    # if url:
-    #     print(f"Using env override for {env_name}: {url}")
-    #     return url
-
-    # check per-service override in profile
-    service_override = profile.current().get("overrides", {}).get(service_name, {})
+
+def _lazy_init(service_name: str, hdc: str = None, region: str = None):  # make it deferred so no need to initialize profile
+    if region and hdc:
+        raise Exception("region and hdc cannot be specified at the same time.")
+
+    profile_data = profile.current()
+
+    service_override = _get_service_override(service_name)
     service_override_url = service_override.get("url")
     if service_override_url:
-        print(f"Using per-service override for {service_name}: {service_override_url}")
-        return service_override_url
-    return profile.current().hcs.url
+        log.debug(f"Using per-service override for {service_name}: {service_override_url}")
+        url = service_override_url
+    else:
+        if hdc:
+            # prod only
+            hdc = hdc.upper()
+            if hdc == "US":
+                url = profile_data.hcs.url
+            elif hdc == "EU":
+                url = profile_data.hcs.url.replace("cloud-sg-us", "cloud-sg-eu")
+            elif hdc == "JP":
+                url = profile_data.hcs.url.replace("cloud-sg-us", "cloud-sg-jp")
+            else:
+                raise CtxpException(f"Invalid HDC name: {hdc}. Supported HDC names: US, EU, JP.")
+        elif region:
+            url = _get_region_url(region)
+            if not url:
+                panic("Missing profile property: hcs.regions")
+        else:
+            url = profile_data.hcs.url
+    url = url.rstrip("/") + "/" + service_name
+
+    # TODO
+    client_id = service_override.get("client-id", None)
+    if not client_id:
+        client_id = service_override.get("clientId", None)
+    client_secret = service_override.get("client-secret", None)
+    if not client_secret:
+        client_secret = service_override.get("clientSecret", None)
+    api_token = service_override.get("api-token", None)
+    if not api_token:
+        api_token = service_override.get("apiToken", None)
+    provider = service_override.get("provider", "vmwarecsp")
+    if provider == "vmwarecsp":
+        token_url = profile_data.csp.url
+    elif provider == "auth-service":
+        token_url = profile_data.auth["tokenUrl"]
+    else:
+        raise CtxpException(f"Unknown provider: {provider}. Supported providers: vmwarecsp, auth-service.")
+
+    if client_id:
+        if not client_secret:
+            panic(f"Client ID is specified but missing clientSecret for service {service_name} in profile override.")
+    else:
+        if client_secret:
+            panic(f"Client secret is specified but missing clientId for service {service_name} in profile override.")
+
+    if client_id:
+        auth_config = {
+            "url": token_url,
+            "client_id": client_id,
+            "client_secret": client_secret,
+        }
+    elif api_token:
+        auth_config = {
+            "url": token_url,
+            "api_token": api_token,
+        }
+    else:
+        auth_config = None
+
+    from hcs_core.sglib import auth
+
+    oauth_client = auth.oauth_client(auth_config=auth_config)
+    return url, oauth_client
 
 
-def _get_region_url(region_name: str):
+def hdc_service_client(service_name: str, hdc: str = None) -> EzClient:
+    """A client for HDC service. Cached per service name, thread-safe, lazy initialization."""
+    if hdc:
+        hdc = hdc.upper()
+        if hdc not in ["US", "EU", "JP"]:
+            panic(f"Invalid HDC name: {hdc}. Supported HDC names: US, EU, JP.")
+    else:
+        hdc = "US"
+
+    k = f"{service_name}#{hdc}"
+    with _client_instance_lock:
+        instance = _caches.get(k)
+        if not instance:
+            instance = EzClient(lazy_init=lambda: _lazy_init(service_name=service_name, hdc=hdc))
+            _caches[k] = instance
+        return instance
+
+
+def _get_region_url(region: str):
     regions = profile.current().hcs.regions
-    if not region_name:
+    if not region:
         return regions[0].url
     for r in regions:
-        if r.name.lower() == region_name.lower():
+        if r.name.lower() == region.lower():
             return r.url
     names = []
     for r in regions:
         names.append(r.name)
-    panic(f"Region not found: {region_name}. Available regions: {names}")
+    panic(f"Region not found: {region}. Available regions: {names}")
 
 
-def regional_service_client(region_name: str, service_name: str):
+def regional_service_client(service_name: str, region: str = None):
     # 'https://dev1b-westus2-cp103a.azcp.horizon.vmware.com/vmhub'
-    url = _get_region_url(region_name)
-    if not url:
-        panic("Missing profile property: hcs.regions")
-    if not url.endswith("/"):
-        url += "/"
-    url += service_name
-    return hcs_client(url)
+    region_names = [r.name.lower() for r in profile.current().hcs.regions]
+    if region:
+        region = region.lower()
+        if region not in region_names:
+            panic(f"Invalid region name: {region}. Available regions: {', '.join(region_names)}")
+    else:
+        region = region_names[0]
 
+    k = f"{service_name}#{region}"
+    with _client_instance_lock:
+        instance = _caches.get(k)
+        if not instance:
+            instance = EzClient(lazy_init=lambda: _lazy_init(service_name=service_name, region=region))
+            _caches[k] = instance
+        return instance
 
-def _with_org_id(url: str, org_id: str):
-    if org_id:
-        if url.find("?") < 0:
-            url += "?"
-        url += "org_id=" + org_id
-    return url
+
+def is_regional_service(service_name: str):
+    regional_services = ["vmhub", "connection-service"]
+    return service_name in regional_services
+
+
+def service_client(service_name: str, region: str = None, hdc: str = None):
+    if is_regional_service(service_name):
+        return regional_service_client(service_name, region)
+    else:
+        return hdc_service_client(service_name, hdc)
 
 
 class default_crud:
@@ -178,7 +261,10 @@ class default_crud:
 
     def wait_for_deleted(self, id: str, org_id: str, timeout: str, fn_is_error: Callable = None):
         name = self._resource_type_name + "/" + id
-        fn_get = lambda: self.get(id, org_id)
+
+        def fn_get():
+            return self.get(id, org_id)
+
         return wait_for_res_deleted(name, fn_get, timeout=timeout, fn_is_error=fn_is_error)
 
     def update(self, id: str, org_id: str, data: dict, **kwargs):
@@ -202,10 +288,13 @@ def wait_for_res_deleted(
     resource_name: str,
     fn_get: Callable,
     timeout: str,
-    polling_interval_seconds: int = 10,
+    polling_interval: int = 10,
     fn_is_error: Callable = None,
 ):
     timeout_seconds = _parse_timeout(timeout)
+    polling_interval_seconds = _parse_timeout(polling_interval)
+    if polling_interval_seconds < 3:
+        polling_interval_seconds = 3
     start = time.time()
     while True:
         t = fn_get()
@@ -224,9 +313,10 @@ def wait_for_res_deleted(
         sleep_seconds = remaining_seconds
         if sleep_seconds > polling_interval_seconds:
             sleep_seconds = polling_interval_seconds
-        exit.sleep(sleep_seconds)
+        time.sleep(sleep_seconds)
 
 
+# flake8: noqa=E731
 def wait_for_res_status(
     resource_name: str,
     fn_get: Callable,
@@ -250,21 +340,25 @@ def wait_for_res_status(
         field_name = get_status
         get_status = lambda t: t[field_name]
     if status_map:
-        if isinstance(status_map["ready"], str):
-            status_map["ready"] = [status_map["ready"]]
-        if isinstance(status_map["transition"], str):
-            status_map["transition"] = [status_map["transition"]]
-        if isinstance(status_map["error"], str):
-            status_map["error"] = [status_map["error"]]
         if is_ready:
             raise CtxpException("Can not specify is_ready when status_map is provided.")
         if is_error:
             raise CtxpException("Can not specify is_error when status_map is provided.")
         if is_transition:
             raise CtxpException("Can not specify is_transition when status_map is provided.")
-        is_ready = lambda s: s in status_map["ready"]
-        is_error = lambda s: s in status_map["error"]
-        is_transition = lambda s: s in status_map["transition"]
+
+        ready_status = status_map["ready"]
+        error_status = status_map["error"]
+        transition_status = status_map["transition"]
+        if isinstance(ready_status, str):
+            ready_status = [ready_status]
+        if isinstance(error_status, str):
+            error_status = [error_status]
+        if isinstance(transition_status, str):
+            transition_status = [transition_status]
+        is_ready = lambda s: s in ready_status
+        is_error = lambda s: error_status is None or s in error_status
+        is_transition = lambda s: transition_status is None or s in transition_status
     else:
         if not is_ready:
             raise CtxpException("Either status_map or is_ready must be specified.")
@@ -272,6 +366,9 @@ def wait_for_res_status(
             raise CtxpException("Either status_map or is_error must be specified.")
         if not is_transition:
             raise CtxpException("Either status_map or is_transition must be specified.")
+        ready_status = None
+        error_status = None
+        transition_status = None
 
     while True:
         t = fn_get()
@@ -281,9 +378,9 @@ def wait_for_res_status(
             raise CtxpException(prefix + "Not found.")
         status = get_status(t)
         if is_error(status):
-            msg = prefix + f"Status error. Actual={status}"
-            if status_map:
-                msg += f", expected={status_map['ready']}"
+            msg = prefix + f"Unexpected terminal state. Actual={status}"
+            if ready_status:
+                msg += f", expected={ready_status}"
             print("-- DUMP START --", file=sys.stderr)
             print(json.dumps(t, indent=4), file=sys.stderr)
             print("-- DUMP END --", file=sys.stderr)
@@ -291,14 +388,15 @@ def wait_for_res_status(
         if is_ready(status):
             return t
         if not is_transition(status):
-            raise CtxpException(
-                prefix + f"Unexpected status: {status}. If this is a transition, add it to status_map['transition']."
-            )
+            raise CtxpException(prefix + f"Unexpected status: {status}. If this is a transition, add it to status_map['transition'].")
 
         now = time.time()
         remaining_seconds = timeout_seconds - (now - start)
         if remaining_seconds < 1:
-            raise TimeoutError(prefix + "Timeout.")
+            msg = prefix + f"Timeout. Current: {status}."
+            if ready_status:
+                msg += f" Expected: {ready_status}."
+            raise TimeoutError(msg)
         sleep_seconds = remaining_seconds
         if sleep_seconds > polling_interval_seconds:
             sleep_seconds = polling_interval_seconds