kaqing 1.77.0__py3-none-any.whl → 2.0.171__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,
+                )

{walker → adam}/sso/authenticator.py

@@ -5,7 +5,11 @@ from .idp_login import IdpLogin
 
 class Authenticator:
     @abstractmethod
-    def authenticate(self, idp_uri: str, app_host: str, username: str, password: str) -> IdpLogin:
+    def name(self) -> str:
+        return None
+
+    @abstractmethod
+    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

@@ -0,0 +1,170 @@
+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 adam.utils import debug
+from .idp_login import IdpLogin
+from adam.config import Config
+
+class AdException(Exception):
+    pass
+
+class AdAuthenticator(Authenticator):
+    def name(self) -> str:
+        return 'ActiveDirectory'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(AdAuthenticator, cls).__new__(cls)
+
+        return cls.instance
+
+    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)
+        state_token = params.get('state', [''])[0]
+        redirect_url = params.get('redirect_uri', [''])[0]
+
+        session = requests.Session()
+        r = session.get(idp_uri)
+        debug(f'{r.status_code} {idp_uri}')
+
+        config = self.validate_and_return_config(r)
+
+        groups = re.match(r'(https://.*?/.*?)/.*', idp_uri)
+        if not groups:
+            raise AdException('Incorrect idp_uri configuration.')
+
+        login_uri = f'{groups[1]}/login'
+        body = {
+            'login': username,
+            'passwd': password,
+            'ctx': config['sCtx'],
+            'hpgrequestid': config['sessionId'],
+            'flowToken': config['sFT']
+        }
+        r = session.post(login_uri, data=body, headers={
+            'Content-Type': 'application/x-www-form-urlencoded'
+        })
+        debug(f'{r.status_code} {login_uri}')
+
+        config = self.validate_and_return_config(r)
+
+        groups = re.match(r'(https://.*?)/.*', idp_uri)
+        if not groups:
+            raise AdException('Incorrect idp_uri configuration.')
+
+        kmsi_uri = f'{groups[1]}/kmsi'
+        body = {
+            'ctx': config['sCtx'],
+            'hpgrequestid': config['sessionId'],
+            'flowToken': config['sFT'],
+        }
+        r = session.post(kmsi_uri, data=body, headers={
+            'Content-Type': 'application/x-www-form-urlencoded'
+        })
+        debug(f'{r.status_code} {kmsi_uri}')
+
+        if (config := self.extract_config_object(r.text)):
+            if 'sErrorCode' in config and config['sErrorCode'] == '50058':
+                raise AdException('Invalid username/password.')
+            elif 'strServiceExceptionMessage' in config:
+                raise AdException(config['strServiceExceptionMessage'])
+            else:
+                Log.log_to_file(config)
+                raise AdException('Unknown err.')
+
+        id_token = self.extract(r.text, r'.*name=\"id_token\" value=\"(.*?)\".*')
+        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)
+        whitelisted = self.whitelisted_members()
+
+        for role in roles:
+            if role in whitelisted:
+                return IdpLogin(redirect_url, id_token, state_token, username, idp_uri=idp_uri, id_token_obj=parsed, session=session)
+
+        contact = Config().get('idps.ad.contact', 'Please contact ted.tran@c3.ai.')
+        raise AdException(f'{username} is not whitelisted. {contact}')
+
+    def validate_and_return_config(self, r: requests.Response):
+        if r.status_code < 200 or r.status_code >= 300:
+            debug(r.text)
+
+            return None
+
+        return self.extract_config_object(r.text)
+
+    def extract_config_object(self, text: str):
+        for line in text.split('\n'):
+            groups = re.match(r'.*\$Config=\s*(\{.*)', line)
+            if groups:
+                js = groups[1].replace(';', '')
+                config = json.loads(js)
+
+                return config
+
+        return None
+
+    def whitelisted_members(self) -> list[str]:
+        members_f = Config().get('idps.ad.whitelist-file', '/kaqing/members')
+        try:
+            with open(members_f, 'r') as file:
+                lines = file.readlines()
+            lines = [line.strip() for line in lines]
+
+            def is_non_comment(line: str):
+                return not line.startswith('#')
+
+            lines = list(filter(is_non_comment, lines))
+
+            return [line.split('#')[0].strip(' ') for line in lines]
+        except FileNotFoundError:
+            pass
+
+        return []
+
+    def parse_id_token(self, id_token: str) -> IdToken:
+        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:
+            debug(traceback.format_exc())
+
+        return None

{walker → adam}/sso/authn_okta.py

@@ -3,20 +3,27 @@ import jwt
 import requests
 from urllib.parse import urlparse, parse_qs, unquote
 
-from walker.sso.authenticator import Authenticator
+from adam.sso.authenticator import Authenticator
+from adam.sso.id_token import IdToken
 
 from .idp_login import IdpLogin
-from walker.config import Config
-from walker.utils import log2
+from adam.config import Config
+from adam.utils import debug, log2
+
+class OktaException(Exception):
+    pass
 
 class OktaAuthenticator(Authenticator):
+    def name(self) -> str:
+        return 'Okta'
+
     # the singleton pattern
     def __new__(cls, *args, **kwargs):
         if not hasattr(cls, 'instance'): cls.instance = super(OktaAuthenticator, cls).__new__(cls)
 
         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)
@@ -25,7 +32,7 @@ class OktaAuthenticator(Authenticator):
 
         okta_host = parsed_url.hostname
 
-        url = f"https://{okta_host}/api/v1/authn"
+        authn_uri = f"https://{okta_host}/api/v1/authn"
         payload = {
             "username": username,
             "password": password,
@@ -41,20 +48,18 @@ class OktaAuthenticator(Authenticator):
         }
 
         session = requests.Session()
-        response = session.post(url, headers=headers, data=json.dumps(payload))
-        if Config().is_debug():
-            log2(f'{response.status_code} {url}')
+        response = session.post(authn_uri, headers=headers, data=json.dumps(payload))
+        debug(f'{response.status_code} {authn_uri}')
         auth_response = response.json()
 
         if 'sessionToken' not in auth_response:
-            return None
+            raise OktaException('Invalid username/password.')
 
         session_token = auth_response['sessionToken']
 
         url = f'{idp_uri}&sessionToken={session_token}'
         r = session.get(url)
-        if Config().is_debug():
-            log2(f'{r.status_code} {url}')
+        debug(f'{r.status_code} {url}')
 
         id_token = OktaAuthenticator().extract(r.text, r'.*name=\"id_token\" value=\"(.*?)\".*')
         if not id_token:
@@ -64,28 +69,32 @@ class OktaAuthenticator(Authenticator):
             else:
                 log2('id_token not found\n' + r.text)
 
-            return None
+            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):
-            if group not in OktaAuthenticator.get_groups(okta_host, id_token):
+            parsed = OktaAuthenticator.parse_id_token(okta_host, id_token)
+            if group not in parsed.groups:
                 tks = group.split('/')
                 group = tks[len(tks) - 1]
                 log2(f'{username} is not a member of {group}.')
 
-                return None
+                raise OktaException("You are not part of admin group.")
 
-        return IdpLogin(redirect_url, id_token, state_token, username, session=session)
+        return IdpLogin(redirect_url, id_token, state_token, username, idp_uri=idp_uri, id_token_obj=parsed, session=session)
 
-    def get_groups(idp_host, id_token) -> list[str]:
-        groups: list[str] = []
+    def parse_id_token(idp_host, id_token) -> IdToken:
+        data: dict[str, any] = []
 
         if not jwt.algorithms.has_crypto:
             log2("No crypto support for JWT, please install the cryptography dependency")
 
-            return groups
+            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)
@@ -103,8 +112,16 @@ class OktaAuthenticator(Authenticator):
                 },
             )
 
-            return data['groups']
+            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:
             pass
 
-        return groups
+        return None

adam/sso/cred_cache.py

@@ -0,0 +1,60 @@
+import os
+from pathlib import Path
+import traceback
+from dotenv import load_dotenv
+
+from adam.config import Config
+from adam.utils import debug
+from adam.utils_k8s.kube_context import KubeContext
+
+class CredCache:
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(CredCache, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self):
+        if not hasattr(self, 'env_f'):
+            self.dir = f'{Path.home()}/.kaqing'
+            self.env_f = f'{self.dir}/.credentials'
+            # immutable - cannot reload with different file content
+            load_dotenv(dotenv_path=self.env_f)
+
+            self.overrides: dict[str, str] = {}
+
+    def get_username(self):
+        return self.overrides['IDP_USERNAME'] if 'IDP_USERNAME' in self.overrides else self.get('IDP_USERNAME')
+
+    def get_password(self):
+        return self.overrides['IDP_PASSWORD'] if 'IDP_PASSWORD' in self.overrides else self.get('IDP_PASSWORD')
+
+    def get(self, key: str) -> str:
+        return os.getenv(key)
+
+    def cache(self, username: str, password: str = None):
+        if os.path.exists(self.env_f):
+            with open(self.env_f, 'w') as file:
+                try:
+                    file.truncate()
+                except:
+                    debug(traceback.format_exc())
+
+        updated = []
+        updated.append(f'IDP_USERNAME={username}')
+        if not KubeContext.in_cluster() and password:
+            # do not store password to the .credentials file when in Kubernetes pod
+            updated.append(f'IDP_PASSWORD={password}')
+
+        if updated:
+            if not os.path.exists(self.env_f):
+                os.makedirs(self.dir, exist_ok=True)
+            with open(self.env_f, 'w') as file:
+                file.write('\n'.join(updated))
+
+            if username:
+                self.overrides['IDP_USERNAME'] = username
+            if password:
+                self.overrides['IDP_PASSWORD'] = password
+
+            debug(f'Cached username: {username}, password: {password}, try load: {self.get_username()}')

adam/sso/id_token.py

@@ -0,0 +1,23 @@
+class IdToken:
+    def __init__(self, data: dict[str, any], email: str, username: str, groups: list[str] = [], iat: int = 0, nbf: int = 0, exp: int = 0):
+        self.data = data
+        self.email = email
+        self.username = username
+        self.groups = groups
+        self.iat = iat
+        self.nbf = nbf
+        self.exp = exp
+
+    def from_dict(j: dict[str, any]):
+        return IdToken(
+            j['data'],
+            j['email'],
+            j['username'],
+            groups=j['groups'] if 'groups' in j else None,
+            iat=j['iat'] if 'iat' in j else None,
+            nbf=j['nbf'] if 'nbf' in j else None,
+            exp=j['exp'] if 'exp' in j else None
+        )
+
+    def to_dict(self):
+        return self.__dict__

adam/sso/idp.py

@@ -0,0 +1,143 @@
+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 debug, log
+
+T = TypeVar('T')
+
+class Idp:
+    ctrl_c_entered = False
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(Idp, cls).__new__(cls)
+
+        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, 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:
+                log(f'{session.idp_host()} login: {username}')
+
+            while not username or Idp.ctrl_c_entered:
+                if Idp.ctrl_c_entered:
+                    Idp.ctrl_c_entered = False
+
+                default_user: str = None
+                if use_cached_creds:
+                    default_user = CredCache().get_username()
+                    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)
+
+                    if forced:
+                        username = default_user
+                    else:
+                        username = input(f'{session.idp_host()} login(default {default_user}): ') or default_user
+                else:
+                    username = input(f'{session.idp_host()} login: ')
+
+            session2: IdpSession = IdpSession.create(username, app_host, app_host)
+            if session.idp_host() != session2.idp_host():
+                session = session2
+
+                log(f'Switching to {session.idp_host()}...')
+                log()
+                log(f'{session.idp_host()} login: {username}')
+
+            password = None
+            while password == None or Idp.ctrl_c_entered: # exit the while loop even if password is empty string
+                if Idp.ctrl_c_entered:
+                    Idp.ctrl_c_entered = False
+
+                default_pass = CredCache().get_password() if use_cached_creds else None
+                if default_pass:
+                    if forced:
+                        password = default_pass
+                    else:
+                        password = Idp.with_no_ican(lambda: getpass.getpass(f'Password(default ********): ') or default_pass)
+                else:
+                    password = Idp.with_no_ican(lambda: getpass.getpass(f'Password: '))
+
+            if username and password:
+                # 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:
+            debug(traceback.format_exc())
+            pass
+
+        return None

adam/sso/idp_login.py

@@ -0,0 +1,50 @@
+import base64
+import json
+from urllib.parse import parse_qs, urlparse
+import requests
+
+from adam.sso.id_token import IdToken
+
+class IdpLogin:
+    def __init__(self, app_login_url: str, id_token: str, state: str, user: str = None, idp_uri: str = None, id_token_obj: IdToken = None, session: requests.Session = None):
+        self.app_login_url = app_login_url
+        self.id_token = id_token
+        self.state = state
+        self.user = user
+        self.idp_uri = idp_uri
+        self.id_token_obj = id_token_obj
+        self.session = session
+
+    def deser(idp_token: str):
+        j = json.loads(base64.b64decode(idp_token.encode('utf-8')))
+
+        return IdpLogin(
+            j['r'],
+            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']) if 'id_token_obj' in j else None)
+
+    def ser(self):
+        return base64.b64encode(json.dumps({
+            'r': self.app_login_url,
+            'id_token': self.id_token,
+            'state': self.state,
+            'idp_uri': self.idp_uri,
+            '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):
+        parsed_url = urlparse(idp_uri)
+        query_string = parsed_url.query
+        params = parse_qs(query_string)
+        state_token = params.get('state', [''])[0]
+        redirect_url = params.get('redirect_uri', [''])[0]
+
+        return IdpLogin(app_login_url=redirect_url, id_token=None, state=state_token, idp_uri=idp_uri)
+
+    def shell_user(self):
+        if not self.user:
+            return None
+
+        return self.user.split('@')[0].replace('.', '')