Merge pull request #13 from MauriceBrg/ft-fix-prefix-bug-add-logout-h…

…elper

Fixes #12 and adds support for logout functionality

dash_cognito_auth/cognito_oauth.py

@@ -1,10 +1,14 @@
+from urllib.parse import quote
+
 from oauthlib.oauth2.rfc6749.errors import InvalidGrantError, TokenExpiredError
 from dash import Dash
 from flask import (
     redirect,
+    request,
     url_for,
     Response,
     session,
+    make_response,
 )
 from .cognito import make_cognito_blueprint, cognito
 
@@ -16,7 +20,14 @@ class CognitoOAuth(Auth):
     Wraps a Dash App and adds Cognito based OAuth2 authentication.
     """
 
-    def __init__(self, app: Dash, domain: str, region=None, additional_scopes=None):
+    def __init__(
+        self,
+        app: Dash,
+        domain: str,
+        region=None,
+        additional_scopes=None,
+        logout_url: str = None,
+    ):
         """
         Wrap a Dash App with Cognito authentication.
 
@@ -43,11 +54,21 @@ def __init__(self, app: Dash, domain: str, region=None, additional_scopes=None):
             AWS region of the User Pool. Mandatory if domain is NOT a custom domain, by default None
         additional_scopes : Additional OAuth Scopes to request, optional
             By default openid, email, and profile are requested - default value: None
+        logout_url : str, optional
+            Add a URL to the app that logs out the user when accessed via HTTP GET.
+            The URL automatically respects path prefixes, i.e. if your app is hosted
+            at example.com/some/prefix and you set logout_url to "logout", the actual
+            URL will be example.com/some/prefix/logout. By default, no logout URL is
+            added and you will have to create your own.
         """
         super().__init__(app)
+
+        dash_base_path = app.get_relative_path("")
+
         cognito_bp = make_cognito_blueprint(
             domain=domain,
             region=region,
+            redirect_url=dash_base_path,
             scope=[
                 "openid",
                 "email",
@@ -56,10 +77,36 @@ def __init__(self, app: Dash, domain: str, region=None, additional_scopes=None):
             + (additional_scopes if additional_scopes else []),
         )
 
-        dash_base_path = app.get_relative_path("")
-
         app.server.register_blueprint(cognito_bp, url_prefix=f"{dash_base_path}/login")
 
+        if logout_url is not None:
+            logout_url = (
+                dash_base_path.removesuffix("/") + "/" + logout_url.removeprefix("/")
+            )
+
+            cognito_hostname = (
+                f"{domain}.auth.{region}.amazoncognito.com"
+                if region is not None
+                else domain
+            )
+
+            @app.server.route(logout_url)
+            def handle_logout():
+
+                post_logout_redirect = (
+                    request.host_url.removesuffix("/") + dash_base_path
+                )
+                cognito_logout_url = (
+                    f"https://{cognito_hostname}/logout?"
+                    + f"client_id={cognito_bp.client_id}&logout_uri={quote(post_logout_redirect)}"
+                )
+
+                response = make_response(redirect(cognito_logout_url))
+
+                # Invalidate the session cookie
+                response.set_cookie("session", "empty", max_age=-3600)
+                return response
+
     def is_authorized(self):
         if not cognito.authorized:
             # send to cognito login

tests/test_end_to_end.py

@@ -45,6 +45,7 @@ def end_to_end_app() -> CognitoOAuth:
         dash_app,
         domain=os.environ["COGNITO_DOMAIN"],
         region=os.environ["COGNITO_REGION"],
+        logout_url="logout",
     )
     auth.app.server.config["COGNITO_OAUTH_CLIENT_ID"] = os.environ[
         "COGNITO_OAUTH_CLIENT_ID"
@@ -70,6 +71,8 @@ def test_end_to_end(end_to_end_app: CognitoOAuth):
     - Follow the redirect to the authorization endpoint
     - Follow the redirect to the app home page (logged in)
     - Check the /session-info endpoint to verify the correct user is logged in
+    - Call the /logout endpoint to end the current session
+    - Check a call to the homepage redirects us to the local cognito endpoint
     """
 
     # Arrange
@@ -121,3 +124,14 @@ def test_end_to_end(end_to_end_app: CognitoOAuth):
     # Verify that the logged in users' email matches the one from the env
     session_info_response = client.get("/session-info")
     assert session_info_response.json["email"] == os.environ["COGNITO_EMAIL"]
+
+    # Log out
+    logout_response = client.get("/logout")
+    assert logout_response.status_code == HTTPStatus.FOUND
+    assert "/logout" in logout_response.location
+
+    # Since we're not longer logged in, we should be redirected to the local
+    # Cognito endpoint.
+    homepage_response = client.get("/")
+    assert homepage_response.status_code == HTTPStatus.FOUND
+    assert homepage_response.location == "/login/cognito"

tests/test_end_to_end_with_path_prefix.py

@@ -0,0 +1,145 @@
+"""
+Integration test that authenticates against a real user pool.
+
+Naturally these are a bit sensitive to the way the Cognito UI is implemented.
+
+This is an almost exact copy of test_end_to_end.py with the main difference
+being that the Dash app isn't hosted at the root path (/) but under a prefix
+(/some/prefix).
+"""
+
+# pylint: disable=W0621
+
+import os
+
+from http import HTTPStatus
+
+import requests
+import pytest
+
+from bs4 import BeautifulSoup
+from dash import Dash, html
+from flask import Flask, session
+
+from dash_cognito_auth import CognitoOAuth
+
+# For our end to end test we don't have HTTPS
+os.environ["OAUTHLIB_INSECURE_TRANSPORT"] = "1"
+
+
+@pytest.fixture
+def end_to_end_app_with_prefix() -> CognitoOAuth:
+    """
+    Small dash app that's wrapped with the CognitoOAuth and offers a /session-info
+    endpoint which returns the currently logged in user.
+    """
+
+    name = "end-to-end"
+
+    dash_app = Dash(name, server=Flask(name), url_base_pathname="/some/prefix/")
+    dash_app.layout = html.H1("Hello World")
+    dash_app.server.config.update(
+        {
+            "TESTING": True,
+        }
+    )
+    dash_app.server.secret_key = "just_a_test"
+
+    auth = CognitoOAuth(
+        dash_app,
+        domain=os.environ["COGNITO_DOMAIN"],
+        region=os.environ["COGNITO_REGION"],
+        logout_url="logout",
+    )
+    auth.app.server.config["COGNITO_OAUTH_CLIENT_ID"] = os.environ[
+        "COGNITO_OAUTH_CLIENT_ID"
+    ]
+    auth.app.server.config["COGNITO_OAUTH_CLIENT_SECRET"] = os.environ[
+        "COGNITO_OAUTH_CLIENT_SECRET"
+    ]
+
+    @dash_app.server.route("/some/prefix/session-info")
+    def session_info():
+        return {"email": session["email"]}
+
+    return auth
+
+
+def test_end_to_end_with_path_prefix(end_to_end_app_with_prefix: CognitoOAuth):
+    """
+
+    End to end test with a Dash app that's hosted unter /some/prefix/
+
+    - Request the local webapp
+    - Follow the redirect to the local authorization endpoint
+    - Follow the redirect to the Cognito UI
+    - Parse the Cognito UI
+    - Log in to Cognito
+    - Follow the redirect to the authorization endpoint
+    - Follow the redirect to the app home page (logged in)
+    - Check the /session-info endpoint to verify the correct user is logged in
+    - Call the /some/prefix/logout endpoint to end the current session
+    - Check a call to the homepage redirects us to the local cognito endpoint
+    """
+
+    # Arrange
+    server: Flask = end_to_end_app_with_prefix.app.server
+    client = server.test_client()
+    s = requests.session()
+
+    # Act + Assert
+
+    # We're not authenticated, we should be redirected to the local cognito endpoint.
+    redirect_to_local_cognito = client.get("/some/prefix/")
+
+    # Redirect to the Cognito Login UI
+    redirect_to_cognito_ui = client.get(redirect_to_local_cognito.location)
+
+    # Get Cognito Login page, extract tokens
+    cognito_login_ui = s.get(redirect_to_cognito_ui.location, timeout=5)
+    ui_soup = BeautifulSoup(cognito_login_ui.text, features="html.parser")
+
+    ui_soup.select_one('input[name="_csrf"]')
+    csrf_token = ui_soup.select_one('input[name="_csrf"]')["value"]
+    cognito_asf_data = ui_soup.select_one('input[name="cognitoAsfData"]')
+    login_url = (
+        redirect_to_cognito_ui.location.split(".com/")[0]
+        + ".com"
+        + ui_soup.select_one('form[name="cognitoSignInForm"]')["action"]
+    )
+
+    # Login and catch redirect response
+    login_response = s.post(
+        url=login_url,
+        data={
+            "_csrf": csrf_token,
+            "username": os.environ["COGNITO_USER_NAME"],
+            "password": os.environ["COGNITO_PASSWORD"],
+            "cognitoAsfData": cognito_asf_data,
+        },
+        timeout=5,
+        allow_redirects=False,
+    )
+
+    # Use Cognito tokens to log in
+    post_cognito_auth_redirect = client.get(login_response.headers["location"])
+
+    # We should now be redirected to the home page which will be displayed
+    assert post_cognito_auth_redirect.location == "/some/prefix/"
+    home_page_with_auth = client.get(post_cognito_auth_redirect.location)
+    assert home_page_with_auth.status_code == HTTPStatus.OK
+
+    # Verify that the logged in users' email matches the one from the env
+    session_info_response = client.get("/some/prefix/session-info")
+    assert session_info_response.json["email"] == os.environ["COGNITO_EMAIL"]
+
+    # Log out
+    logout_response = client.get("/some/prefix/logout")
+    assert logout_response.status_code == HTTPStatus.FOUND
+    assert "/logout" in logout_response.location
+
+    # Since we're not longer logged in, we should be redirected to the local
+    # Cognito endpoint.
+    homepage_response = client.get("/some/prefix/")
+    assert homepage_response.status_code == HTTPStatus.FOUND
+    assert homepage_response.location == "/some/prefix/login/cognito"