kaqing 1.77.0__py3-none-any.whl → 2.0.171__py3-none-any.whl

adam/sso/idp_session.py

@@ -0,0 +1,55 @@
+import os
+from urllib.parse import urlparse
+
+from adam.sso import authn_okta
+from adam.sso.authenticator import Authenticator
+from adam.sso.authn_ad import AdAuthenticator
+from adam.sso.sso_config import SsoConfig
+from .idp_login import IdpLogin
+
+class IdpNotSupportedException(Exception):
+    pass
+
+class IdpSession:
+    def create(email: str, client_id: str, app_host: str, idp_uri: str = None) -> 'IdpSession':
+        if not idp_uri:
+            idp_uri = SsoConfig().find_idp_uri(email, client_id, app_host)
+
+        return IdpSession(idp_uri, IdpSession._authenticator(idp_uri))
+
+    def __init__(self, idp_uri: str, authenticator: Authenticator):
+        self.idp_uri = idp_uri
+        self.authenticator = authenticator
+
+    def idp_host(self):
+        return urlparse(self.idp_uri).hostname
+
+    def _authenticator(idp_uri) -> Authenticator:
+        idp = urlparse(idp_uri).hostname
+
+        if 'okta' in idp.lower():
+            return authn_okta.OktaAuthenticator()
+        elif 'microsoftonline' in idp.lower():
+            return AdAuthenticator()
+
+        raise IdpNotSupportedException(f'{idp} is not supported; only okta and ad are supported.')
+
+    def login_from_env_var(self) -> IdpLogin:
+        if idp_token := os.getenv('IDP_TOKEN'):
+            return self.login_from_token(idp_token)
+
+        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
+
+        l0.state = l1.state
+
+        return l0
+
+    def get_idp_login(self) -> IdpLogin:
+        return IdpLogin.create_from_idp_uri(self.idp_uri)

adam/sso/sso_config.py

@@ -0,0 +1,63 @@
+import base64
+import os
+import re
+
+from adam.config import Config
+
+class IdpConfigurationException(Exception):
+    pass
+
+class SsoConfig:
+    DEFAULT_IDP_URI = 'https://c3energy.okta.com/oauth2/v1/authorize?response_type=id_token&response_mode=form_post&client_id={client_id}&scope=openid+email+profile+groups&redirect_uri=https%3A%2F%2F{host}%2Fc3%2Fc3%2Foidc%2Flogin&nonce={nonce}&state=EMPTY'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(SsoConfig, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self):
+        if not hasattr(self, 'config'):
+            self.config = Config().get('idps', {
+                'ad': {
+                    'email-pattern': '.*@c3.ai',
+                    'uri': 'https://login.microsoftonline.com/53ad779a-93e7-485c-ba20-ac8290d7252b/oauth2/v2.0/authorize?response_type=id_token&response_mode=form_post&client_id=00ff94a8-6b0a-4715-98e0-95490012d818&scope=openid+email+profile&redirect_uri=https%3A%2F%2Fplat.c3ci.cloud%2Fc3%2Fc3%2Foidc%2Flogin&nonce={nonce}&state=EMPTY',
+                    'contact': 'Please contact ted.tran@c3.ai.'
+                },
+                'okta': {
+                    'default': True,
+                    'email-pattern': '.*@c3iot.com',
+                    'uri': SsoConfig.DEFAULT_IDP_URI
+                }
+            })
+
+    def find_idp_uri(self, user_email: str, client_id: str, app_host: str) -> str:
+        default: str = SsoConfig.DEFAULT_IDP_URI
+        for k, v in self.config.items():
+            if 'default' in v and v['default']:
+                default = v['uri']
+                break
+
+        if not default:
+            raise IdpConfigurationException('Please configure default idp.')
+
+        def rpl(uri: str):
+            return uri.replace('{client_id}', client_id).replace('{host}', app_host).replace('{nonce}', SsoConfig.generate_oauth_nonce())
+
+        if not user_email:
+            return rpl(default)
+
+        for idp_name, conf in self.config.items():
+            if 'email-pattern' in conf:
+                groups = re.match(conf['email-pattern'], user_email)
+                if groups:
+                    return rpl(conf['uri'])
+
+        return rpl(default)
+
+    def generate_oauth_nonce():
+        random_bytes = os.urandom(32)
+        nonce = base64.urlsafe_b64encode(random_bytes).decode('utf-8')
+        nonce = nonce.rstrip('=')
+
+        return nonce