kaqing 1.98.15__py3-none-any.whl → 2.0.145__py3-none-any.whl

adam/sql/term_completer.py

@@ -0,0 +1,76 @@
+from typing import Callable, Iterable, List, Mapping, Optional, Pattern, Union
+
+from prompt_toolkit.completion import CompleteEvent, Completion, WordCompleter
+from prompt_toolkit.document import Document
+from prompt_toolkit.formatted_text import AnyFormattedText
+
+__all__ = [
+    "TermCompleter",
+]
+
+class TermCompleter(WordCompleter):
+    def __init__(
+        self,
+        words: Union[List[str], Callable[[], List[str]]],
+        ignore_case: bool = False,
+        display_dict: Optional[Mapping[str, AnyFormattedText]] = None,
+        meta_dict: Optional[Mapping[str, AnyFormattedText]] = None,
+        WORD: bool = False,
+        sentence: bool = False,
+        match_middle: bool = False,
+        pattern: Optional[Pattern[str]] = None,
+    ) -> None:
+        super().__init__(words, ignore_case, display_dict, meta_dict, WORD, sentence, match_middle, pattern)
+
+    def __str__(self):
+        return ','.join(self.words)
+
+    def get_completions(
+        self, document: Document, complete_event: CompleteEvent
+    ) -> Iterable[Completion]:
+        # Get list of words.
+        words = self.words
+        if callable(words):
+            words = words()
+
+        # Get word/text before cursor.
+        if self.sentence:
+            word_before_cursor = document.text_before_cursor
+        else:
+            word_before_cursor = document.get_word_before_cursor(
+                WORD=self.WORD, pattern=self.pattern
+            )
+
+        if self.ignore_case:
+            word_before_cursor = word_before_cursor.lower()
+
+        def word_matches(word: str) -> bool:
+            """True when the word before the cursor matches."""
+            if self.ignore_case:
+                word = word.lower()
+
+            if self.match_middle:
+                return word_before_cursor in word
+            else:
+                return word.startswith(word_before_cursor)
+
+        for a in words:
+            if word_before_cursor in ['(', ',', '=', 'in', "',"]:
+                display = self.display_dict.get(a, a)
+                display_meta = self.meta_dict.get(a, "")
+                yield Completion(
+                    a,
+                    0,
+                    display=display,
+                    display_meta=display_meta,
+                )
+
+            if word_matches(a):
+                display = self.display_dict.get(a, a)
+                display_meta = self.meta_dict.get(a, "")
+                yield Completion(
+                    a,
+                    -len(word_before_cursor),
+                    display=display,
+                    display_meta=display_meta,
+                )

adam/sso/authenticator.py

@@ -9,7 +9,7 @@ class Authenticator:
         return None
 
     @abstractmethod
-    def authenticate(self, idp_uri: str, app_host: str, username: str, password: str) -> IdpLogin:
+    def authenticate(self, idp_uri: str, app_host: str, username: str, password: str, verify: bool = True) -> IdpLogin:
         pass
 
     def extract(self, form: str, pattern: re.Pattern):

adam/sso/authn_ad.py

@@ -1,13 +1,13 @@
-import base64
 import json
 import re
+import traceback
+import jwt
 import requests
 from urllib.parse import urlparse, parse_qs
 
 from adam.log import Log
 from adam.sso.authenticator import Authenticator
 from adam.sso.id_token import IdToken
-
 from .idp_login import IdpLogin
 from adam.config import Config
 
@@ -24,7 +24,7 @@ class AdAuthenticator(Authenticator):
 
         return cls.instance
 
-    def authenticate(self, idp_uri: str, app_host: str, username: str, password: str) -> IdpLogin:
+    def authenticate(self, idp_uri: str, app_host: str, username: str, password: str, verify: bool) -> IdpLogin:
         parsed_url = urlparse(idp_uri)
         query_string = parsed_url.query
         params = parse_qs(query_string)
@@ -84,6 +84,9 @@ class AdAuthenticator(Authenticator):
         if not id_token:
             raise AdException('Invalid username/password.')
 
+        if not verify:
+            return IdpLogin(redirect_url, id_token, state_token, username, idp_uri=idp_uri, id_token_obj=None, session=session)
+
         parsed = self.parse_id_token(id_token)
         roles = parsed.groups
         roles.append(username)
@@ -134,56 +137,33 @@ class AdAuthenticator(Authenticator):
         return []
 
     def parse_id_token(self, id_token: str) -> IdToken:
-        def decode_jwt_part(encoded_part):
-            missing_padding = len(encoded_part) % 4
-            if missing_padding:
-                encoded_part += '=' * (4 - missing_padding)
-            decoded_bytes = base64.urlsafe_b64decode(encoded_part)
-            return json.loads(decoded_bytes.decode('utf-8'))
-
-        parts = id_token.split('.')
-        # header = decode_jwt_part(parts[0])
-        data = decode_jwt_part(parts[1])
-        # print('SEAN', payload)
-        # {
-        #     'aud': '00ff94a8-6b0a-4715-98e0-95490012d818',
-        #     'iss': 'https://login.microsoftonline.com/53ad779a-93e7-485c-ba20-ac8290d7252b/v2.0',
-        #     'iat': 1756138348,
-        #     'nbf': 1756138348,
-        #     'exp': 1756142248,
-        #     'cc': 'CmCuOndNgXFbP/Vor0kv9fqd32LOv7kSHKStVrGPTXXnlPlSET3g4z23XjVZJW37F2Yy8d45MzZ6xA/XNbYGE3BHYAZFhfDKOp0ZbWZysqa2zqD3lpyXxnlEpzWkFY1SlDgSBWMzLmFpGhIKEFbS9ukejKFBi0QKCcPHh6MiEgoQb7gHorgqqEOuuPDIcSgHATICTkE4AUIJCQDvjjtx391I',
-        #     'email': 'sean.ahn@c3.ai',
-        #     'groups': [
-        #         'e58ef97f-8622-47cc-93cd-0ec8e3df2b4e',
-        #         '290214ec-2eaa-47bb-802c-70c2535bb7e7',
-        #         '55cd2e92-c40d-4646-b837-c6fb6406013b',
-        #         '9b715aa3-ec6c-44ad-be0c-a1d95045526d',
-        #         '19a38c19-e4c3-4d7b-8bb0-4afe7f502a51',
-        #         '5d891ce3-02d8-4d34-9748-9e711a3d54c5',
-        #         '0626cb36-106c-4ab3-adcf-7ee8e3f05584',
-        #         '6dfdfa3e-b225-4a74-a534-1cc963f78e08',
-        #         'bfdcca4e-7212-4e73-9b8a-eecfd3a7797d',
-        #         '6d050845-61c4-4072-acb1-9662d0c9faa0',
-        #         'c4fbb32c-9892-4eb3-a829-b4eaaf71b4ef'
-        #     ],
-        #     'name': 'Sean Ahn',
-        #     'nonce': 'V7DzmHtmhu3X3tzZmC55vMQ2WJtqXnW6wJpYb3Kfud8',
-        #     'oid': '380029a1-0643-4764-b3b8-9bc02630af41',
-        #     'preferred_username': 'sean.ahn@c3.ai',
-        #     'rh': '1.ATcAmnetU-eTXEi6IKyCkNclK6iU_wAKaxVHmOCVSQAS2Bj1AJU3AA.',
-        #     'sid': '007dea79-a65c-7911-5142-8ea1b9faa41a',
-        #     'sub': 'd-iCznXBDo3HEV7UCGalCVFIG47dQl_SFCaQtN2yVQI',
-        #     'tid': '53ad779a-93e7-485c-ba20-ac8290d7252b',
-        #     'uti': 'b7gHorgqqEOuuPDIcSgHAQ',
-        #     'ver': '2.0'
-        # }
-
-        return IdToken(
-            data,
-            data['email'],
-            data['name'],
-            groups=data['groups'] if 'groups' in data else [],
-            iat=data['iat'] if 'iat' in data else 0,
-            nbf=data['nbf'] if 'nbf' in data else 0,
-            exp=data['exp'] if 'exp' in data else 0
-        )
+        jwks_url = Config().get('idps.ad.jwks-uri', '')
+        try:
+            jwks_client = jwt.PyJWKClient(jwks_url, cache_jwk_set=True, lifespan=360)
+            signing_key = jwks_client.get_signing_key_from_jwt(id_token)
+            data = jwt.decode(
+                id_token,
+                signing_key.key,
+                algorithms=["RS256"],
+                options={
+                    "verify_signature": True,
+                    "verify_exp": False,
+                    "verify_nbf": True,
+                    "verify_iat": True,
+                    "verify_aud": False,
+                    "verify_iss": False,
+                },
+            )
+            return IdToken(
+                data,
+                data['email'],
+                data['name'],
+                groups=data['groups'] if 'groups' in data else [],
+                iat=data['iat'] if 'iat' in data else 0,
+                nbf=data['nbf'] if 'nbf' in data else 0,
+                exp=data['exp'] if 'exp' in data else 0
+            )
+        except:
+            Config().debug(traceback.format_exc())
+
+        return None

adam/sso/authn_okta.py

@@ -23,7 +23,7 @@ class OktaAuthenticator(Authenticator):
 
         return cls.instance
 
-    def authenticate(self, idp_uri: str, app_host: str, username: str, password: str) -> IdpLogin:
+    def authenticate(self, idp_uri: str, app_host: str, username: str, password: str, verify: bool) -> IdpLogin:
         parsed_url = urlparse(idp_uri)
         query_string = parsed_url.query
         params = parse_qs(query_string)
@@ -53,7 +53,6 @@ class OktaAuthenticator(Authenticator):
         auth_response = response.json()
 
         if 'sessionToken' not in auth_response:
-            print('password', password, auth_response)
             raise OktaException('Invalid username/password.')
 
         session_token = auth_response['sessionToken']
@@ -72,6 +71,10 @@ class OktaAuthenticator(Authenticator):
 
             raise OktaException('Invalid username/password.')
 
+        if not verify:
+            # just relay id_token, it will be verified by SP
+            return IdpLogin(redirect_url, id_token, state_token, username, idp_uri=idp_uri, id_token_obj=None, session=session)
+
         if group := Config().get('app.login.admin-group', '{host}/C3.ClusterAdmin').replace('{host}', app_host):
             parsed = OktaAuthenticator.parse_id_token(okta_host, id_token)
             if group not in parsed.groups:
@@ -80,34 +83,6 @@ class OktaAuthenticator(Authenticator):
                 log2(f'{username} is not a member of {group}.')
 
                 raise OktaException("You are not part of admin group.")
-            # print(parsed)
-            # {
-            #   'sub': '00u1fii2azgmh46dE1d8',
-            #   'name': 'Sean Ahn',
-            #   'locale': 'US',
-            #   'email': 'sean.ahn@c3iot.com',
-            #   'ver': 1,
-            #   'iss': 'https://c3energy.okta.com',
-            #   'aud': 'azops88.c3.ai',
-            #   'iat': 1756136559,
-            #   'exp': 1756140159,
-            #   'jti': 'ID.fIjNa_A-uYo-mKMqpN_CLPqgQUnDIHeBV4WGa-Q_Qx8',
-            #   'amr': [
-            #     'pwd'
-            #   ],
-            #   'idp': '00onsxj5mmBGYLRURHTH',
-            #   'nonce': 'kRpFTGiyUGg6uWfdCUkS3tGI4wrroCz3MTaeI_wTILE',
-            #   'preferred_username': 'sean.ahn@c3iot.com',
-            #   'given_name': 'Seung',
-            #   'family_name': 'Ahn',
-            #   'zoneinfo': 'America/Los_Angeles',
-            #   'updated_at': 1755201182,
-            #   'email_verified': False,
-            #   'auth_time': 1756136558,
-            #   'groups': [
-            #     'azops88.c3.ai/C3.ClusterAdmin'
-            #   ]
-            # }
 
         return IdpLogin(redirect_url, id_token, state_token, username, idp_uri=idp_uri, id_token_obj=parsed, session=session)
 
@@ -119,8 +94,7 @@ class OktaAuthenticator(Authenticator):
 
             return None
 
-        okta_auth_server = f"https://{idp_host}/oauth2"
-        jwks_url = f"{okta_auth_server}/v1/keys"
+        jwks_url = Config().get('idps.okta.jwks-uri', 'https://c3energy.okta.com/oauth2/v1/keys')
         try:
             jwks_client = jwt.PyJWKClient(jwks_url, cache_jwk_set=True, lifespan=360)
             signing_key = jwks_client.get_signing_key_from_jwt(id_token)

adam/sso/cred_cache.py

@@ -4,7 +4,7 @@ import traceback
 from dotenv import load_dotenv
 
 from adam.config import Config
-from adam.k8s_utils.kube_context import KubeContext
+from adam.utils_k8s.kube_context import KubeContext
 
 class CredCache:
     # the singleton pattern

adam/sso/idp.py

@@ -1,10 +1,23 @@
+import base64
 import getpass
+import os
+import sys
+import termios
+import traceback
+from typing import Callable, TypeVar
+import requests
+from kubernetes import config
+import yaml
+
+from adam.utils_k8s.secrets import Secrets
 
 from .cred_cache import CredCache
 from .idp_session import IdpSession
 from .idp_login import IdpLogin
 from adam.config import Config
-from adam.utils import log
+from adam.utils import log, log2
+
+T = TypeVar('T')
 
 class Idp:
     ctrl_c_entered = False
@@ -15,13 +28,23 @@ class Idp:
 
         return cls.instance
 
-    def login(app_host: str, username: str = None, idp_uri: str = None, forced = False, use_token_from_env = True, use_cached_creds = True) -> IdpLogin:
+    def login(app_host: str, username: str = None, idp_uri: str = None, forced = False, use_token_from_env = True, use_cached_creds = True, verify = True) -> IdpLogin:
         session: IdpSession = IdpSession.create(username, app_host, app_host, idp_uri=idp_uri)
 
         if use_token_from_env:
             if l0 := session.login_from_env_var():
                 return l0
 
+            if port := os.getenv("SERVER_PORT"):
+                token_server = Config().get('app.login.token-server-url', 'http://localhost:{port}').replace('{port}', port)
+                res: requests.Response = requests.get(token_server)
+                if res.status_code == 200 and res.text:
+                    try:
+                        # may fail if the idp token is not complete
+                        return session.login_from_token(res.text)
+                    except:
+                        pass
+
         r: IdpLogin = None
         try:
             if username:
@@ -31,8 +54,13 @@ class Idp:
                 if Idp.ctrl_c_entered:
                     Idp.ctrl_c_entered = False
 
-                default_user = CredCache().get_username() if use_cached_creds else None
-                Config().debug(f'User read from cache: {default_user} with use_cached: {use_cached_creds}')
+                default_user: str = None
+                if use_cached_creds:
+                    default_user = CredCache().get_username()
+                    Config().debug(f'User read from cache: {default_user}')
+
+                if from_env := os.getenv('USERNAME'):
+                    default_user = from_env
                 if default_user and default_user != username:
                     session = IdpSession.create(default_user, app_host, app_host)
 
@@ -61,18 +89,55 @@ class Idp:
                     if forced:
                         password = default_pass
                     else:
-                        password = getpass.getpass(f'Password(default ********): ') or default_pass
+                        password = Idp.with_no_ican(lambda: getpass.getpass(f'Password(default ********): ') or default_pass)
                 else:
-                    password = getpass.getpass(f'Password: ')
+                    password = Idp.with_no_ican(lambda: getpass.getpass(f'Password: '))
 
             if username and password:
-                r = session.authenticator.authenticate(session.idp_uri, app_host, username, password)
-
-                return r
+                # if uploading kubeconfig file fails many times, you will be locked out
+                # kubeconfig file content has first char as tab or length of bigger than 128
+                if password[0] == '\t' or len(password) > Config().get('app.login.password-max-length', 128):
+                    if r := Idp.try_kubeconfig(username, password):
+                        log(f"You're signed in as {username}")
+                        return r
+                else:
+                    if r := session.authenticator.authenticate(session.idp_uri, app_host, username, password, verify=verify):
+                        log(f"You're signed in as {username}")
+                    return r
         finally:
             if r and Config().get('app.login.cache-creds', True):
                 CredCache().cache(username, password)
             elif username and Config().get('app.login.cache-username', True):
                 CredCache().cache(username)
 
+        return None
+
+    def with_no_ican(body: Callable[[], T]) -> T:
+        # override 4096 character limit with OS terminal - stty -noican
+        fd = sys.stdin.fileno()
+        old = termios.tcgetattr(fd)
+        new = termios.tcgetattr(fd)
+        new[3] = new[3] & ~termios.ICANON
+        try:
+            termios.tcsetattr(fd, termios.TCSADRAIN, new)
+            return body()
+        finally:
+            termios.tcsetattr(fd, termios.TCSADRAIN, old)
+
+    def try_kubeconfig(username: str, kubeconfig: str):
+        try:
+            if kubeconfig[0] == '\t':
+                kubeconfig = kubeconfig[1:]
+            kubeconfig_string = base64.b64decode(kubeconfig.encode('ascii') + b'==').decode('utf-8')
+            if kubeconfig_string.startswith('apiVersion: '):
+                kubeconfig_dict = yaml.safe_load(kubeconfig_string)
+                config.kube_config.load_kube_config_from_dict(kubeconfig_dict)
+                # test if you can list Kubernetes secretes with the given kubeconfig file
+                Secrets.list_secrets(os.getenv('NAMESPACE'))
+
+                return IdpLogin(None, None, None, username)
+        except:
+            Config().debug(traceback.format_exc())
+            pass
+
         return None

adam/sso/idp_login.py

@@ -23,7 +23,7 @@ class IdpLogin:
             j['id_token'],
             j['state'],
             idp_uri=j['idp_uri'] if 'idp_uri' in j else None,
-            id_token_obj=IdToken.from_dict(j['id_token_obj']))
+            id_token_obj=IdToken.from_dict(j['id_token_obj']) if 'id_token_obj' in j else None)
 
     def ser(self):
         return base64.b64encode(json.dumps({
@@ -31,7 +31,7 @@ class IdpLogin:
             'id_token': self.id_token,
             'state': self.state,
             'idp_uri': self.idp_uri,
-            'id_token_obj': self.id_token_obj.to_dict()
+            'id_token_obj': self.id_token_obj.to_dict() if self.id_token_obj else None
         }).encode('utf-8')).decode('utf-8')
 
     def create_from_idp_uri(idp_uri: str):

adam/sso/idp_session.py

@@ -36,17 +36,20 @@ class IdpSession:
 
     def login_from_env_var(self) -> IdpLogin:
         if idp_token := os.getenv('IDP_TOKEN'):
-            l0: IdpLogin = IdpLogin.deser(idp_token)
-            l1: IdpLogin = self.get_idp_login()
-            # if l0.app_login_url == l1.app_login_url:
-            if l0.state != 'EMPTY':
-                return l0
+            return self.login_from_token(idp_token)
 
-            l0.state = l1.state
+        return None
 
+    def login_from_token(self, idp_token: str) -> IdpLogin:
+        l0: IdpLogin = IdpLogin.deser(idp_token)
+        l1: IdpLogin = self.get_idp_login()
+        # if l0.app_login_url == l1.app_login_url:
+        if l0.state != 'EMPTY':
             return l0
 
-        return None
+        l0.state = l1.state
+
+        return l0
 
     def get_idp_login(self) -> IdpLogin:
         return IdpLogin.create_from_idp_uri(self.idp_uri)

adam/utils.py

@@ -9,7 +9,8 @@ import os
 from pathlib import Path
 import random
 import string
-from typing import cast
+import threading
+from typing import Callable
 from dateutil import parser
 import subprocess
 import sys
@@ -19,6 +20,8 @@ import yaml
 
 from . import __version__
 
+is_debug_holder = [lambda: False]
+
 def to_tabular(lines: str, header: str = None, dashed_line = False):
     return lines_to_tabular(lines.split('\n'), header, dashed_line)
 
@@ -70,15 +73,21 @@ def epoch(timestamp_string: str):
     return parser.parse(timestamp_string).timestamp()
 
 def log(s = None):
+    if not loggable():
+        return
+
     # want to print empty line for False or empty collection
     if s == None:
         print()
     else:
         click.echo(s)
 
-def log2(s = None):
+def log2(s = None, nl = True):
+    if not loggable():
+        return
+
     if s:
-        click.echo(s, err=True)
+        click.echo(s, err=True, nl=nl)
     else:
         print(file=sys.stderr)
 
@@ -125,6 +134,19 @@ def deep_merge_dicts(dict1, dict2):
             merged_dict[key] = value
     return merged_dict
 
+def deep_sort_dict(d):
+    """
+    Recursively sorts a dictionary by its keys, and any nested lists by their elements.
+    """
+    if not isinstance(d, (dict, list)):
+        return d
+
+    if isinstance(d, dict):
+        return {k: deep_sort_dict(d[k]) for k in sorted(d)}
+
+    if isinstance(d, list):
+        return sorted([deep_sort_dict(item) for item in d])
+
 def get_deep_keys(d, current_path=""):
     """
     Recursively collects all combined keys (paths) from a deep dictionary.
@@ -228,4 +250,63 @@ def copy_config_file(rel_path: str, module: str, suffix: str = '.yaml', show_out
     return path
 
 def idp_token_from_env():
-    return os.getenv('IDP_TOKEN')
+    return os.getenv('IDP_TOKEN')
+
+def is_lambda(func):
+    return callable(func) and hasattr(func, '__name__') and func.__name__ == '<lambda>'
+
+class Ing:
+    state = threading.local()
+
+    def __init__(self, msg: str, suppress_log=False):
+        self.msg = msg
+        self.suppress_log = suppress_log
+        self.nested = False
+
+    def __enter__(self):
+        if not hasattr(Ing.state, 'ing_cnt'):
+            Ing.state.ing_cnt = 0
+
+        try:
+            if not Ing.state.ing_cnt:
+                if not self.suppress_log and not is_debug_holder[0]():
+                    log2(f'{self.msg}...', nl=False)
+
+            return None
+        finally:
+            Ing.state.ing_cnt += 1
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        Ing.state.ing_cnt -= 1
+        if not Ing.state.ing_cnt:
+            if not self.suppress_log and not is_debug_holder[0]():
+                log2(' OK')
+
+        return False
+
+def ing(msg: str, body: Callable[[], None]=None, suppress_log=False):
+    if not body:
+        return Ing(msg, suppress_log=suppress_log)
+
+    r = None
+
+    if not hasattr(Ing.state, 'ing_cnt'):
+        Ing.state.ing_cnt = 0
+
+    if not Ing.state.ing_cnt:
+        if not suppress_log and not is_debug_holder[0]():
+            log2(f'{msg}...', nl=False)
+
+    Ing.state.ing_cnt += 1
+    try:
+        r = body()
+    finally:
+        Ing.state.ing_cnt -= 1
+        if not Ing.state.ing_cnt:
+            if not suppress_log and not is_debug_holder[0]():
+                log2(' OK')
+
+    return r
+
+def loggable():
+    return is_debug_holder[0]() or not hasattr(Ing.state, 'ing_cnt') or not Ing.state.ing_cnt