Add OKTA auth integration

docs_website/docs/integrations/add_auth.md

@@ -23,13 +23,35 @@ We will go through how to setup authentication with OAuth and LDAP starting from
 
 Start by creating an OAuth client with the authentication provider (e.g. [Google](https://developers.google.com/identity/protocols/oauth2), [Okta](https://developer.okta.com/docs/guides/implement-oauth-for-okta/create-oauth-app/)). Make sure "http://localhost:10001/oauth2callback" is entered as allowed redirect uri. Once created, the next step is to change the querybook config by editing `containers/bundled_querybook_config.yaml`. Open that file and enter the following:
 
+#### Generic OAuth
+
 ```yaml
 AUTH_BACKEND: 'app.auth.oauth_auth' # Same as import path when running Python
 OAUTH_CLIENT_ID: '---Redacted---'
 OAUTH_CLIENT_SECRET: '---Redacted---'
 OAUTH_AUTHORIZATION_URL: https://accounts.google.com/o/oauth2/v2/auth
 OAUTH_TOKEN_URL: https://oauth2.googleapis.com/token
 OAUTH_USER_PROFILE: https://openidconnect.googleapis.com/v1/userinfo
+PUBLIC_URL: http://localhost:10001
+```
+
+#### Google
+
+```yaml
+AUTH_BACKEND: 'app.auth.google_auth'
+OAUTH_CLIENT_ID: '---Redacted---'
+OAUTH_CLIENT_SECRET: '---Redacted---'
+PUBLIC_URL: http://localhost:10001
+```
+
+#### Okta
+
+```yaml
+AUTH_BACKEND: 'app.auth.okta_auth'
+OAUTH_CLIENT_ID: '---Redacted---'
+OAUTH_CLIENT_SECRET: '---Redacted---'
+OKTA_BASE_URL: https://[Redacted].okta.com/oauth2
+PUBLIC_URL: http://localhost:10001
 ```
 
 :::caution

package.json

@@ -1,6 +1,6 @@
 {
     "name": "querybook",
-    "version": "2.5.2",
+    "version": "2.5.3",
     "description": "A Big Data Webapp",
     "private": true,
     "scripts": {

querybook/server/app/auth/oauth_auth.py

@@ -61,15 +61,51 @@ def init_app(self, flask_app):
         )
 
     def login(self, request):
-        oauth_url, _ = self.oauth_session.authorization_url(
-            self.oauth_config["authorization_url"]
-        )
+        oauth_url, _ = self._get_authn_url()
         flask_session["next"] = request.path
         return redirect(oauth_url)
 
-    def _parse_user_profile(self, profile_response):
-        user = profile_response.json()["user"]
-        return user["username"], user["email"]
+    def _get_authn_url(self):
+        return self.oauth_session.authorization_url(
+            self.oauth_config["authorization_url"]
+        )
+
+    def oauth_callback(self):
+        LOG.debug("Handling Oauth callback...")
+
+        if request.args.get("error"):
+            return f"<h1>Error: {request.args.get('error')}</h1>"
+
+        code = request.args.get("code")
+
+        try:
+            access_token = self._fetch_access_token(code)
+            username, email = self._get_user_profile(access_token)
+            with DBSession() as session:
+                flask_login.login_user(
+                    AuthUser(self.login_user(username, email, session=session))
+                )
+        except AuthenticationError:
+            abort_unauthorized()
+
+        next_url = "/"
+        if "next" in flask_session:
+            next_url = flask_session["next"]
+            del flask_session["next"]
+
+        return redirect(next_url)
+
+    def _fetch_access_token(self, code):
+        resp = self.oauth_session.fetch_token(
+            token_url=self.oauth_config["token_url"],
+            client_id=self.oauth_config["client_id"],
+            code=code,
+            client_secret=self.oauth_config["client_secret"],
+            cert=certifi.where(),
+        )
+        if resp is None:
+            raise AuthenticationError("Null response, denying access.")
+        return resp["access_token"]
 
     def _get_user_profile(self, access_token):
         resp = requests.get(
@@ -83,6 +119,10 @@ def _get_user_profile(self, access_token):
             )
         return self._parse_user_profile(resp)
 
+    def _parse_user_profile(self, profile_response):
+        user = profile_response.json()["user"]
+        return user["username"], user["email"]
+
     @with_session
     def login_user(self, username, email, session=None):
         user = get_user_by_name(username, session=session)
@@ -92,41 +132,6 @@ def login_user(self, username, email, session=None):
             )
         return user
 
-    def oauth_callback(self):
-        LOG.debug("Handling Oauth callback...")
-        if request.args.get("error"):
-            return f"<h1>Error: {request.args.get('error')}</h1>"
-
-        resp = self.oauth_session.fetch_token(
-            token_url=self.oauth_config["token_url"],
-            client_id=self.oauth_config["client_id"],
-            code=request.args.get("code"),
-            client_secret=self.oauth_config["client_secret"],
-            cert=certifi.where(),
-        )
-
-        try:
-            if resp is None:
-                raise AuthenticationError("Null response, denying access.")
-
-            access_token = resp["access_token"]
-
-            username, email = self._get_user_profile(access_token)
-        except AuthenticationError:
-            abort_unauthorized()
-
-        with DBSession() as session:
-            flask_login.login_user(
-                AuthUser(self.login_user(username, email, session=session))
-            )
-
-        next_url = "/"
-        if "next" in flask_session:
-            next_url = flask_session["next"]
-            del flask_session["next"]
-
-        return redirect(next_url)
-
 
 login_manager = OAuthLoginManager()
 

querybook/server/app/auth/okta_auth.py

@@ -0,0 +1,90 @@
+import certifi
+import requests
+
+from app.auth.oauth_auth import OAuthLoginManager, OAUTH_CALLBACK_PATH
+from env import QuerybookSettings, get_env_config
+from lib.utils.decorators import in_mem_memoized
+from .utils import AuthenticationError
+
+
+class NoopAuth(requests.auth.AuthBase):
+    """
+    This auth doesn't do anything.
+    It only used to override oauthlib's behavior.
+    """
+
+    def __call__(self, r):
+        return r
+
+
+class OktaLoginManager(OAuthLoginManager):
+    def get_okta_urls(self):
+        okta_base_url = get_env_config("OKTA_BASE_URL")
+        authorization_url = f"{okta_base_url}/v1/authorize"
+        token_url = f"{okta_base_url}/v1/token"
+        profile_url = f"{okta_base_url}/v1/userinfo"
+        return authorization_url, token_url, profile_url
+
+    @property
+    @in_mem_memoized()
+    def oauth_config(self):
+        authorization_url, token_url, profile_url = self.get_okta_urls()
+
+        return {
+            "callback_url": "{}{}".format(
+                QuerybookSettings.PUBLIC_URL, OAUTH_CALLBACK_PATH
+            ),
+            "client_id": QuerybookSettings.OAUTH_CLIENT_ID,
+            "client_secret": QuerybookSettings.OAUTH_CLIENT_SECRET,
+            "authorization_url": authorization_url,
+            "token_url": token_url,
+            "profile_url": profile_url,
+            "scope": ["openid", "email"],
+        }
+
+    def _fetch_access_token(self, code):
+        resp = self.oauth_session.fetch_token(
+            token_url=self.oauth_config["token_url"],
+            client_id=self.oauth_config["client_id"],
+            code=code,
+            client_secret=self.oauth_config["client_secret"],
+            cert=certifi.where(),
+            # This Authentication is needed because Okta would throw error
+            # about passing client_secret and client_id in request.header
+            # which is the default behavior of oauthlib
+            auth=NoopAuth(),
+        )
+        if resp is None:
+            raise AuthenticationError("Null response, denying access.")
+        return resp["access_token"]
+
+    def _get_user_profile(self, access_token):
+        resp = requests.get(
+            self.oauth_config["profile_url"],
+            headers={"Authorization": "Bearer {}".format(access_token)},
+        )
+        if not resp or resp.status_code != 200:
+            raise AuthenticationError(
+                "Failed to fetch user profile, status ({0})".format(
+                    resp.status if resp else "None"
+                )
+            )
+        return self._parse_user_profile(resp)
+
+    def _parse_user_profile(self, resp):
+        user = resp.json()
+        username = user["email"].split("@")[0]
+        return username, user["email"]
+
+
+login_manager = OktaLoginManager()
+
+ignore_paths = [OAUTH_CALLBACK_PATH]
+
+
+def init_app(app):
+    login_manager.init_app(app)
+
+
+def login(request):
+    return login_manager.login(request)

querybook/server/env.py

@@ -14,7 +14,7 @@ class MissingConfigException(Exception):
     pass
 
 
-def get_dh_config(name, optional=True):
+def get_env_config(name, optional=True):
     found = True
     if name in os.environ:
         val = os.environ.get(name)
@@ -36,52 +36,52 @@ def get_dh_config(name, optional=True):
 class QuerybookSettings(object):
     # Core
     PRODUCTION = os.environ.get("production", "false") == "true"
-    PUBLIC_URL = get_dh_config("PUBLIC_URL")
-    FLASK_SECRET_KEY = get_dh_config("FLASK_SECRET_KEY", optional=False)
+    PUBLIC_URL = get_env_config("PUBLIC_URL")
+    FLASK_SECRET_KEY = get_env_config("FLASK_SECRET_KEY", optional=False)
 
     # Celery
-    REDIS_URL = get_dh_config("REDIS_URL", optional=False)
+    REDIS_URL = get_env_config("REDIS_URL", optional=False)
 
     # Search
-    ELASTICSEARCH_HOST = get_dh_config("ELASTICSEARCH_HOST", optional=False)
-    ELASTICSEARCH_CONNECTION_TYPE = get_dh_config("ELASTICSEARCH_CONNECTION_TYPE")
+    ELASTICSEARCH_HOST = get_env_config("ELASTICSEARCH_HOST", optional=False)
+    ELASTICSEARCH_CONNECTION_TYPE = get_env_config("ELASTICSEARCH_CONNECTION_TYPE")
 
     # Database
-    DATABASE_CONN = get_dh_config("DATABASE_CONN", optional=False)
-    DATABASE_POOL_SIZE = int(get_dh_config("DATABASE_POOL_SIZE"))
-    DATABASE_POOL_RECYCLE = int(get_dh_config("DATABASE_POOL_RECYCLE"))
+    DATABASE_CONN = get_env_config("DATABASE_CONN", optional=False)
+    DATABASE_POOL_SIZE = int(get_env_config("DATABASE_POOL_SIZE"))
+    DATABASE_POOL_RECYCLE = int(get_env_config("DATABASE_POOL_RECYCLE"))
 
     # Communications
-    EMAILER_CONN = get_dh_config("EMAILER_CONN")
-    QUERYBOOK_SLACK_TOKEN = get_dh_config("QUERYBOOK_SLACK_TOKEN")
-    QUERYBOOK_EMAIL_ADDRESS = get_dh_config("QUERYBOOK_EMAIL_ADDRESS")
+    EMAILER_CONN = get_env_config("EMAILER_CONN")
+    QUERYBOOK_SLACK_TOKEN = get_env_config("QUERYBOOK_SLACK_TOKEN")
+    QUERYBOOK_EMAIL_ADDRESS = get_env_config("QUERYBOOK_EMAIL_ADDRESS")
 
     # Authentication
-    AUTH_BACKEND = get_dh_config("AUTH_BACKEND")
-    LOGS_OUT_AFTER = int(get_dh_config("LOGS_OUT_AFTER"))
+    AUTH_BACKEND = get_env_config("AUTH_BACKEND")
+    LOGS_OUT_AFTER = int(get_env_config("LOGS_OUT_AFTER"))
 
-    OAUTH_CLIENT_ID = get_dh_config("OAUTH_CLIENT_ID")
-    OAUTH_CLIENT_SECRET = get_dh_config("OAUTH_CLIENT_SECRET")
-    OAUTH_AUTHORIZATION_URL = get_dh_config("OAUTH_AUTHORIZATION_URL")
-    OAUTH_TOKEN_URL = get_dh_config("OAUTH_TOKEN_URL")
-    OAUTH_USER_PROFILE = get_dh_config("OAUTH_USER_PROFILE")
+    OAUTH_CLIENT_ID = get_env_config("OAUTH_CLIENT_ID")
+    OAUTH_CLIENT_SECRET = get_env_config("OAUTH_CLIENT_SECRET")
+    OAUTH_AUTHORIZATION_URL = get_env_config("OAUTH_AUTHORIZATION_URL")
+    OAUTH_TOKEN_URL = get_env_config("OAUTH_TOKEN_URL")
+    OAUTH_USER_PROFILE = get_env_config("OAUTH_USER_PROFILE")
 
-    LDAP_CONN = get_dh_config("LDAP_CONN")
-    LDAP_USER_DN = get_dh_config("LDAP_USER_DN")
+    LDAP_CONN = get_env_config("LDAP_CONN")
+    LDAP_USER_DN = get_env_config("LDAP_USER_DN")
 
     # Result Store
-    RESULT_STORE_TYPE = get_dh_config("RESULT_STORE_TYPE")
+    RESULT_STORE_TYPE = get_env_config("RESULT_STORE_TYPE")
 
-    STORE_BUCKET_NAME = get_dh_config("STORE_BUCKET_NAME")
-    STORE_PATH_PREFIX = get_dh_config("STORE_PATH_PREFIX")
-    STORE_MIN_UPLOAD_CHUNK_SIZE = int(get_dh_config("STORE_MIN_UPLOAD_CHUNK_SIZE"))
-    STORE_MAX_UPLOAD_CHUNK_NUM = int(get_dh_config("STORE_MAX_UPLOAD_CHUNK_NUM"))
-    STORE_MAX_READ_SIZE = int(get_dh_config("STORE_MAX_READ_SIZE"))
-    STORE_READ_SIZE = int(get_dh_config("STORE_READ_SIZE"))
+    STORE_BUCKET_NAME = get_env_config("STORE_BUCKET_NAME")
+    STORE_PATH_PREFIX = get_env_config("STORE_PATH_PREFIX")
+    STORE_MIN_UPLOAD_CHUNK_SIZE = int(get_env_config("STORE_MIN_UPLOAD_CHUNK_SIZE"))
+    STORE_MAX_UPLOAD_CHUNK_NUM = int(get_env_config("STORE_MAX_UPLOAD_CHUNK_NUM"))
+    STORE_MAX_READ_SIZE = int(get_env_config("STORE_MAX_READ_SIZE"))
+    STORE_READ_SIZE = int(get_env_config("STORE_READ_SIZE"))
 
-    DB_MAX_UPLOAD_SIZE = int(get_dh_config("DB_MAX_UPLOAD_SIZE"))
+    DB_MAX_UPLOAD_SIZE = int(get_env_config("DB_MAX_UPLOAD_SIZE"))
 
-    GOOGLE_CREDS = json.loads(get_dh_config("GOOGLE_CREDS") or "null")
+    GOOGLE_CREDS = json.loads(get_env_config("GOOGLE_CREDS") or "null")
 
     # Logging
-    LOG_LOCATION = get_dh_config("LOG_LOCATION")
+    LOG_LOCATION = get_env_config("LOG_LOCATION")

querybook/server/lib/utils/decorators.py

@@ -42,6 +42,7 @@ def expired(self):
 
 def in_mem_memoized(ttl_secs=None):
     """Memoizes the results of the function.
+       Note params change are ignored in memo.
 
     Args:
         func: A function that will have its results memoized.