Add oauth2 authorization code flow support

CONTRIBUTING.md

@@ -9,6 +9,7 @@
   - [Data input](#data-input)
     - [Deattaching data from its original source](#deattaching-data-from-its-original-source)
   - [Custom ad-hoc plugins](#custom-ad-hoc-plugins)
+  - [OAuth 2.0 Authorization Code flow](#oauth-2.0-authorization-code-flow)
 - [Run tests](#run-tests)
 - [Build documentation](#build-documentation)
 
@@ -531,6 +532,35 @@ applies if the custom plugin is saved in a package with submodule(s),
           ...
 ```
 
+### OAuth 2.0 Authorization Code flow
+
+It is possible to use OAuth 2.0 Authorization Code flow to secure a `webviz` application. 
+In order to do so, add `oauth2` attribute in a custom plugin with a boolean `True` value. 
+The following is an example.
+
+```python
+class OurCustomPlugin(WebvizPluginABC):
+
+    def __init__(self):
+        super().__init__()
+        self.use_oauth2 = True
+
+    @property
+    def oauth2(self):
+        return self.use_oauth2
+```
+
+Information related to the application for OAuth 2.0 has to be provided in environment 
+variables. These environment variables are `WEBVIZ_TENANT_ID`, `WEBVIZ_CLIENT_ID`, 
+`WEBVIZ_CLIENT_SECRET`, `WEBVIZ_SCOPE`.
+
+The values can be found in the Azure AD configuration page. Short explanation of these environment variables:
+
+- `WEBVIZ_TENANT_ID`: The organization's Azure tenant ID (Equinor has exactly one tenant ID).
+- `WEBVIZ_CLIENT_ID`: ID of the Webviz Azure AD app.
+- `WEBVIZ_CLIENT_SECRET`: Webviz Azure AD app's client secret.
+- `WEBVIZ_SCOPE`: The API permission for this Webviz Azure AD app.
+
 ## Run tests
 
 To run tests it is necessary to first install the [selenium chrome driver](https://github.com/SeleniumHQ/selenium/wiki/ChromeDriver).

setup.py

@@ -80,6 +80,7 @@ def get_long_description() -> str:
         "flask-talisman>=0.6",
         "jinja2>=2.10",
         "markdown>=3.0",
+        "msal>=1.5.0",
         "pandas>=1.0",
         "pyarrow>=0.16",
         "pyyaml>=5.1",

webviz_config/__init__.py

@@ -11,6 +11,7 @@
 from ._is_reload_process import is_reload_process
 from ._plugin_abc import WebvizPluginABC, EncodedFile, ZipFileMember
 from ._shared_settings_subscriptions import SHARED_SETTINGS_SUBSCRIPTIONS
+from ._oauth2 import Oauth2
 
 try:
     __version__ = version("webviz-config")

webviz_config/_localhost_token.py

@@ -4,6 +4,7 @@
 import flask
 
 from ._is_reload_process import is_reload_process
+from ._oauth2 import Oauth2
 
 
 class LocalhostToken:
@@ -29,9 +30,10 @@ class LocalhostToken:
     two different localhost applications running simultaneously do not interfere.
     """
 
-    def __init__(self, app: flask.app.Flask, port: int):
+    def __init__(self, app: flask.app.Flask, port: int, oauth2: Oauth2 = None):
         self._app = app
         self._port = port
+        self._oauth2 = oauth2
 
         if not is_reload_process():
             # One time token (per run) user has to provide
@@ -63,6 +65,7 @@ def set_request_decorators(self) -> None:
         # pylint: disable=inconsistent-return-statements
         @self._app.before_request
         def _check_for_ott_or_cookie():  # type: ignore[no-untyped-def]
+
             if not self._ott_validated and self._ott == flask.request.args.get("ott"):
                 self._ott_validated = True
                 flask.g.set_cookie_token = True
@@ -72,6 +75,16 @@ def _check_for_ott_or_cookie():  # type: ignore[no-untyped-def]
                 f"cookie_token_{self._port}"
             ):
                 self._ott_validated = True
+
+                if self._oauth2:
+                    # The session of the request does not contain access token, redirect to /login
+                    is_redirected, redirect_url = self._oauth2.is_empty_token()
+                    if is_redirected:
+                        return flask.redirect(redirect_url)
+
+                    # The session contains access token, check (and set) its expiration date
+                    self._oauth2.check_and_set_token_expiry()
+
             else:
                 flask.abort(401)
 

webviz_config/_oauth2.py

@@ -0,0 +1,162 @@
+import os
+import datetime
+from typing import Tuple
+
+import msal
+import flask
+
+
+class Oauth2:
+    """Oauth2 authorization Code grant flow"""
+
+    def __init__(self, app: flask.app.Flask):
+        self._app = app
+
+        # Azure AD app registration info (currently the values are taken from environment variables)
+        self._tenant_id = os.environ["WEBVIZ_TENANT_ID"]
+        self._client_id = os.environ["WEBVIZ_CLIENT_ID"]
+        self._client_secret = os.environ["WEBVIZ_CLIENT_SECRET"]
+        self._scope = os.environ["WEBVIZ_SCOPE"]
+
+        # Initiate msal
+        self._msal_app = msal.ConfidentialClientApplication(
+            client_id=self._client_id,
+            client_credential=self._client_secret,
+            authority=f"https://login.microsoftonline.com/{self._tenant_id}",
+        )
+        self._accounts = self._msal_app.get_accounts()
+
+        # Initiate oauth2 endpoints
+        self.set_oauth2_endpoints()
+
+    def set_oauth2_endpoints(self) -> None:
+        """/login and /auth-return endpoints are added for Oauth2 authorization
+        code flow.
+
+        At the end of the flow, a session cookie containing a valid access token
+        and its expiration date is created. This flask session object can be
+        accessed from Webviz plugin.
+
+        To get the access token: flask.session.get("access_token")
+        To get the expiration date: flask.session.get("expiration_date")
+
+        An Azure AD application should be registered, and the following environment
+        variables should be set: WEBVIZ_TENANT_ID, WEBVIZ_CLIENT_ID,
+        WEBVIZ_CLIENT_SECRET, WEBVIZ_SCOPE.
+
+        WEBVIZ_SESSION_SECRET_KEY environment variable is also needed for signing the
+        session cookie.
+        """
+
+        @self._app.route("/login")
+        def _login_controller():  # type: ignore[no-untyped-def]
+            redirect_uri = get_auth_redirect_uri(flask.request.url_root)
+
+            # First leg of Oauth2 authorization code flow
+            auth_url = self._msal_app.get_authorization_request_url(
+                scopes=[self._scope], redirect_uri=redirect_uri
+            )
+            return flask.redirect(auth_url)
+
+        @self._app.route("/auth-return")
+        def _auth_return_controller():  # type: ignore[no-untyped-def]
+            redirect_uri = get_auth_redirect_uri(flask.request.url_root)
+            returned_query_params = flask.request.args
+
+            # There is an error from the first leg of Oauth2 authorization code flow
+            if "error" in returned_query_params:
+                error_description = returned_query_params.get("error_description")
+                print("Error description:", error_description)
+                redirect_error_uri = flask.url_for("error_controller")
+                return flask.redirect(redirect_error_uri)
+
+            code = returned_query_params.get("code")
+
+            # Second leg of Oauth2 authorization code flow
+            tokens_result = self._msal_app.acquire_token_by_authorization_code(
+                code=code, scopes=[self._scope], redirect_uri=redirect_uri
+            )
+            expires_in = tokens_result.get("expires_in")
+            expiration_date = datetime.datetime.now() + datetime.timedelta(
+                seconds=expires_in - 60
+            )
+            print("Access token expiration date:", expiration_date)
+
+            # Set expiration date in the session
+            flask.session["expiration_date"] = expiration_date
+
+            # Set access token in the session
+            flask.session["access_token"] = tokens_result.get("access_token")
+
+            return flask.redirect(flask.request.url_root)
+
+        @self._app.route("/error")
+        def _error_controller():  # type: ignore[no-untyped-def]
+            return "Error"
+
+    def set_oauth2_before_request_decorator(self) -> None:
+        """Check access token existence in session cookie before every request.
+        If it does not exist, the browser is redirected to /login endpoint.
+
+        If access token exists, its expiration date is checked in session cookie.
+        If the current date exceeds its expiration date, a new access token is
+        retrieved and set in the session cookie. A new expiration date is also
+        set in the session cookie.
+        """
+
+        # pylint: disable=inconsistent-return-statements
+        @self._app.before_request
+        def _check_access_token():  # type: ignore[no-untyped-def]
+            # The session of the request does not contain access token, redirect to /login
+            is_redirected, redirect_url = self.is_empty_token()
+            if is_redirected:
+                return flask.redirect(redirect_url)
+
+            # The session contains access token, check its expiration date
+            self.check_and_set_token_expiry()
+
+    @staticmethod
+    def is_empty_token() -> Tuple[bool, str]:
+        if (
+            not flask.session.get("access_token")
+            and flask.request.path != "/login"
+            and flask.request.path != "/auth-return"
+        ):
+            login_uri = get_login_uri(flask.request.url_root)
+            return True, login_uri
+
+        return False, ""
+
+    def check_and_set_token_expiry(self) -> None:
+        if flask.session.get("access_token"):
+            expiration_date = flask.session.get("expiration_date")
+            current_date = datetime.datetime.now()
+            if current_date > expiration_date:
+                # Access token has expired
+                print("Access token has expired.")
+                if not self._accounts:
+                    self._accounts = self._msal_app.get_accounts()
+                renewed_tokens_result = self._msal_app.acquire_token_silent(
+                    scopes=[self._scope], account=self._accounts[0]
+                )
+                expires_in = renewed_tokens_result.get("expires_in")
+                new_expiration_date = datetime.datetime.now() + datetime.timedelta(
+                    seconds=expires_in - 60
+                )
+                print("New access token expiration date:", new_expiration_date)
+
+                # Set new expiration date in the session
+                flask.session["expiration_date"] = new_expiration_date
+
+                # Set new access token in the session
+                flask.session["access_token"] = renewed_tokens_result.get(
+                    "access_token"
+                )
+
+
+def get_login_uri(url_root: str) -> str:
+    return url_root + "login"
+
+
+def get_auth_redirect_uri(url_root: str) -> str:
+    return url_root + "auth-return"

webviz_config/templates/webviz_template.py.jinja2

@@ -39,6 +39,10 @@ app = dash.Dash(
     ],
 )
 app.logger.setLevel(logging.WARNING)
+
+# For signing session cookies
+app.server.secret_key = webviz_config.LocalhostToken.generate_token()
+
 server = app.server
 
 app.title = "{{ title }}"
@@ -66,13 +70,15 @@ app._deprecated_webviz_settings = {
 CACHE.init_app(server)
 
 theme.adjust_csp({"script-src": app.csp_hashes()}, append=True)
-Talisman(server, content_security_policy=theme.csp, feature_policy=theme.feature_policy, force_https=False)
+Talisman(server, content_security_policy=theme.csp, feature_policy=theme.feature_policy, force_https=False, session_cookie_secure=False)
 
 WEBVIZ_STORAGE.use_storage = {{portable}}
 WEBVIZ_STORAGE.storage_folder = Path(__file__).resolve().parent / "webviz_storage"
 
 WEBVIZ_ASSETS.portable = {{ portable }}
 
+use_oauth2 = False
+
 if {{ not portable }} and not webviz_config.is_reload_process():
     # When Dash/Flask is started on localhost with hot module reload activated,
     # we do not want the main process to call expensive component functions in
@@ -82,16 +88,17 @@ if {{ not portable }} and not webviz_config.is_reload_process():
 else:
     page_content = {}
     {% for page in pages %}
-    page_content["{{page.id}}"] = [
-       {% for content in page.content -%}
-       {%- if content is string -%}
-         dcc.Markdown(r"""{{ content }}""")
-       {%- else -%}
-         webviz_config.plugins.{{ content._call_signature[0] }}.{{ content._call_signature[1] }}
-       {%- endif -%}
-       {{- "" if loop.last else ","}}
-       {% endfor -%}
-    ]
+    page_content["{{page.id}}"] = []
+    {% for content in page.content -%}
+    {% if content is string %}
+    page_content["{{page.id}}"].append(dcc.Markdown(r"""{{ content }}"""))
+    {% else %}
+    plugin_content = webviz_config.plugins.{{ content._call_signature[0] }}
+    if not use_oauth2:
+        use_oauth2 = plugin_content.oauth2 if hasattr(plugin_content, "oauth2") else use_oauth2
+    page_content["{{page.id}}"].append(plugin_content.{{ content._call_signature[1] }})
+    {% endif %}
+    {% endfor %}
     {% endfor %}
     app.layout = html.Div(
         className="layoutWrapper", 
@@ -124,6 +131,7 @@ else:
             ])]),
         html.Div(className="pageContent", id="page-content")])
 
+oauth2 = webviz_config.Oauth2(app.server) if use_oauth2 else None
 
 @app.callback(dash.dependencies.Output("page-content", "children"),
               dash.dependencies.Input("url", "pathname"))
@@ -144,7 +152,7 @@ if __name__ == "__main__":
 
     port = webviz_config.utils.get_available_port(preferred_port=5000)
 
-    token = webviz_config.LocalhostToken(app.server, port).one_time_token
+    token = webviz_config.LocalhostToken(app.server, port, oauth2).one_time_token
     webviz_config.utils.LocalhostOpenBrowser(port, token)
 
     webviz_config.utils.silence_flask_startup()
@@ -161,3 +169,7 @@ if __name__ == "__main__":
         dev_tools_silence_routes_logging=False,
       {% endif %}
     )
+else:
+    # This will be applied if not running on localhost
+    if use_oauth2:
+        oauth2.set_oauth2_before_request_decorator()