Merge pull request #11 from MauriceBrg/feat-custom-domain-support

Add support for custom domains and clean up

README.md

@@ -88,8 +88,4 @@ COGNITO_DOMAIN=<just-the-prefix>
 COGNITO_REGION=<aws-region-of-the-cognito-userpool>
 COGNITO_OAUTH_CLIENT_ID=<app-client-id>
 COGNITO_OAUTH_CLIENT_SECRET=<app-client-secret>
-```
-
-## Known Limitations
-
-- Fully Custom Cognito Domains aren't supported at the moment
+```

dash_cognito_auth/auth.py

@@ -1,4 +1,3 @@
-from __future__ import absolute_import
 from abc import ABCMeta, abstractmethod
 from six import iteritems, add_metaclass
 
@@ -7,24 +6,25 @@
 class Auth(object):
     def __init__(self, app):
         self.app = app
-        self._index_view_name = app.config['routes_pathname_prefix']
+        self._index_view_name = app.config["routes_pathname_prefix"]
         self._overwrite_index()
         self._protect_views()
-        self._index_view_name = app.config['routes_pathname_prefix']
+        self._index_view_name = app.config["routes_pathname_prefix"]
 
     def _overwrite_index(self):
         original_index = self.app.server.view_functions[self._index_view_name]
 
-        self.app.server.view_functions[self._index_view_name] = \
-            self.index_auth_wrapper(original_index)
+        self.app.server.view_functions[self._index_view_name] = self.index_auth_wrapper(
+            original_index
+        )
 
     def _protect_views(self):
         # require auth wrapper for all views
-        for view_name, view_method in iteritems(
-                self.app.server.view_functions):
+        for view_name, view_method in iteritems(self.app.server.view_functions):
             if view_name != self._index_view_name:
-                self.app.server.view_functions[view_name] = \
-                    self.auth_wrapper(view_method)
+                self.app.server.view_functions[view_name] = self.auth_wrapper(
+                    view_method
+                )
 
     @abstractmethod
     def is_authorized(self):

dash_cognito_auth/cognito.py

@@ -1,5 +1,3 @@
-from __future__ import unicode_literals
-
 from flask_dance.consumer import OAuth2ConsumerBlueprint
 from flask.globals import LocalProxy
 from flask import g
@@ -19,7 +17,7 @@ def make_cognito_blueprint(
     session_class=None,
     storage=None,
     domain=None,
-    region="eu-west-1",
+    region=None,
 ):
     """
     Make a blueprint for authenticating with Cognito using OAuth 2. This requires
@@ -49,21 +47,33 @@ def make_cognito_blueprint(
                 :class:`~flask_dance.consumer.storage.session.SessionStorage`.
         domain (str): The domain configured in Cognito
         region (str): The region of AWS
-            Defaults to ``eu-west-1``.
 
     :rtype: :class:`~flask_dance.consumer.OAuth2ConsumerBlueprint`
     :returns: A :ref:`blueprint <flask:blueprints>` to attach to your Flask app.
     """
+
+    # There are more sophisticated checks, but for our purposes it should
+    # strike a balance between accuracy and readability. The value of domain
+    # is either just a prefix in Cognito or a FQDN.
+    custom_domain = "." in domain
+
+    if not custom_domain and region is None:
+        raise ValueError("The region parameter must be set if 'domain' is not a FQDN.")
+
+    hostname = (
+        f"{domain}.auth.{region}.amazoncognito.com" if region is not None else domain
+    )
+
     scope = scope or ["openid", "email", "phone", "profile"]
     cognito_bp = OAuth2ConsumerBlueprint(
         "cognito",
         __name__,
         client_id=client_id,
         client_secret=client_secret,
         scope=scope,
-        base_url=f"https://{domain}.auth.{region}.amazoncognito.com",
-        authorization_url=f"https://{domain}.auth.{region}.amazoncognito.com/oauth2/authorize",
-        token_url=f"https://{domain}.auth.{region}.amazoncognito.com/oauth2/token",
+        base_url=f"https://{hostname}",
+        authorization_url=f"https://{hostname}/oauth2/authorize",
+        token_url=f"https://{hostname}/oauth2/token",
         redirect_url=redirect_url,
         redirect_to=redirect_to,
         login_url=login_url,

dash_cognito_auth/cognito_oauth.py

@@ -16,8 +16,35 @@ class CognitoOAuth(Auth):
     Wraps a Dash App and adds Cognito based OAuth2 authentication.
     """
 
-    def __init__(self, app: Dash, domain, region, additional_scopes=None):
-        super(CognitoOAuth, self).__init__(app)
+    def __init__(self, app: Dash, domain: str, region=None, additional_scopes=None):
+        """
+        Wrap a Dash App with Cognito authentication.
+
+        The app needs two configuration options to work:
+
+        COGNITO_OAUTH_CLIENT_ID -> Client-ID of the Cognito App Client
+        COGNITO_OAUTH_CLIENT_SECRET -> Secret of the Cognito App Client
+
+        Can be set like this:
+
+        app.server.config["COGNITO_OAUTH_CLIENT_ID"] = "something"
+        app.server.config["COGNITO_OAUTH_CLIENT_SECRET"] = "something"
+
+        ---
+
+        Parameters
+        ----------
+        app : Dash
+            The app to add authentication to.
+        domain : str
+            Either the domain prefix of the User Pool domain if hosted by Cognito
+            or the FQDN of your custom domain, e.g. authentication.example.com
+        region : str, optional
+            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
+        """
+        super().__init__(app)
         cognito_bp = make_cognito_blueprint(
             domain=domain,
             region=region,

setup.py

@@ -5,7 +5,7 @@
 setup(
     name="dash-cognito-auth",
     description="Dash Cognito Auth",
-    long_description=open("README.md", "rt").read().strip(),
+    long_description=open("README.md", "rt", encoding="utf-8").read().strip(),
     long_description_content_type="text/markdown",
     author="Frank Spijkerman",
     author_email="[email protected]",

tests/conftest.py

@@ -3,6 +3,7 @@
 """
 
 # pylint: disable=W0621
+from typing import Iterator
 from unittest.mock import patch
 
 import pytest
@@ -91,7 +92,23 @@ def prefixed_app_with_auth(app_with_url_prefix) -> CognitoOAuth:
 
 
 @pytest.fixture
-def authorized_app(app) -> CognitoOAuth:
+def app_with_auth_and_cognito_custom_domain(app) -> CognitoOAuth:
+    """
+    Dash App wrapped with Cognito Authentication
+    - Domain name authentication.example.com
+    - App Client Id: testclient
+    - App Client Secret: testsecret
+    """
+
+    auth = CognitoOAuth(app, domain="authentication.example.com")
+    auth.app.server.config["COGNITO_OAUTH_CLIENT_ID"] = "testclient"
+    auth.app.server.config["COGNITO_OAUTH_CLIENT_SCRET"] = "testsecret"
+
+    return auth
+
+
+@pytest.fixture
+def authorized_app(app) -> Iterator[CognitoOAuth]:
     """
     App with Cognito Based authentication that bypasses the authentication/authorization
     part, i.e. replaced is_authorized and the authorized endpoint.

tests/test_auth_flows.py

@@ -118,5 +118,42 @@ def test_that_cognito_authorized_response_is_accepted(authorized_app: CognitoOAu
     ), "Should redirect to root after authorization"
 
 
-# Test Cognito Response parsing
-# TODO: Test Logout
+def test_that_cognito_handler_redirects_to_user_pool_with_custom_domain(
+    app_with_auth_and_cognito_custom_domain: CognitoOAuth,
+):
+    """
+    If we're not logged in, the Cognito handler should redirect the request to the
+    Cognito User pool so the client can login an retrieve a JWT.
+
+    We test that all the required scopes etc. are present in the redirect uri
+    and it's formed according to the expected pattern.
+    In this scenario, we additionally ensure that a custom domain works, in our case
+    it should be authentication.example.com
+    """
+
+    # Arrange
+    flask_server: Flask = app_with_auth_and_cognito_custom_domain.app.server
+    client = flask_server.test_client()
+
+    # Act
+    response = client.get("/login/cognito")
+
+    # Assert
+    assert response.status_code == HTTPStatus.FOUND, "We expect a redirect"
+
+    redirect_url = response.headers.get("Location")
+    parsed = urlparse(redirect_url)
+
+    assert parsed.scheme == "https"
+    assert parsed.hostname == "authentication.example.com"  # Custom Domain Name
+    assert parsed.path == "/oauth2/authorize"
+
+    parsed_qs = parse_qs(parsed.query, strict_parsing=True)
+    assert "openid" in parsed_qs["scope"][0]
+    assert "email" in parsed_qs["scope"][0]
+    assert "profile" in parsed_qs["scope"][0]
+
+    assert parsed_qs["response_type"][0] == "code"
+    assert parsed_qs["redirect_uri"][0] == "http://localhost/login/cognito/authorized"
+    assert parsed_qs["client_id"][0] == "testclient"
+    assert "state" in parsed_qs

tests/test_end_to_end.py

@@ -4,6 +4,8 @@
 Naturally these are a bit sensitive to the way the Cognito UI is implemented.
 """
 
+# pylint: disable=W0621
+
 import os
 
 from http import HTTPStatus

tests/test_import.py

@@ -2,6 +2,7 @@
 Test Dash Cognito Auth.
 """
 
+import pytest
 from dash_cognito_auth import CognitoOAuth
 
 
@@ -11,3 +12,16 @@ def test_init(app):
     auth = CognitoOAuth(app, domain="test", region="eu-west-1")
 
     assert auth.app is app
+
+
+def test_that_init_raises_an_exception_if_cognito_domain_and_region_is_missing(app):
+    """
+    Initializing the app with a Cognito Domain (non-FQDN) and no Region should
+    raise a ValueError.
+    """
+
+    # Arrange
+
+    # Act + Assert
+    with pytest.raises(ValueError):
+        CognitoOAuth(app, "non-fqdn")