add pkce auth (#156)

Author: RomanBachaloSigmaSoftware

app/docusign/ds_client.py

@@ -1,10 +1,14 @@
 import uuid
-from os import path
+from os import path, urandom
+import hashlib
+import base64
+import secrets
 
 
 import requests
-from flask import current_app as app, url_for, redirect, render_template, request, session
+from flask import current_app as app, url_for, redirect, render_template, request, session, redirect
 from flask_oauthlib.client import OAuth
+from requests_oauthlib import OAuth2Session
 from docusign_esign import ApiClient
 from docusign_esign.client.api_exception import ApiException
 
@@ -48,7 +52,10 @@ class DSClient:
     @classmethod
     def _init(cls, auth_type, api):
         if auth_type == "code_grant":
-            cls._auth_code_grant(api)
+            if session.get("pkce_failed", False):
+                cls._auth_code_grant(api)
+            else:
+                cls._pkce_auth(api)
         elif auth_type == "jwt":
             cls._jwt_auth(api)
 
@@ -92,6 +99,29 @@ def _auth_code_grant(cls, api):
             access_token_method="POST"
         )
 
+    @classmethod
+    def _pkce_auth(cls, api):
+        """Authorize with the Authorization Code Grant - OAuth 2.0 flow"""
+        use_scopes = []
+
+        if api == "Rooms":
+            use_scopes.extend(ROOMS_SCOPES)
+        elif api == "Click":
+            use_scopes.extend(CLICK_SCOPES)
+        elif api == "Admin":
+            use_scopes.extend(ADMIN_SCOPES)
+        elif api == "Maestro":
+            use_scopes.extend(MAESTRO_SCOPES)
+        elif api == "WebForms":
+            use_scopes.extend(WEBFORMS_SCOPES)
+        else:
+            use_scopes.extend(SCOPES)
+        # remove duplicate scopes
+        use_scopes = list(set(use_scopes))
+
+        redirect_uri = DS_CONFIG["app_url"] + url_for("ds.ds_callback")
+        cls.ds_app = OAuth2Session(DS_CONFIG["ds_client_id"], redirect_uri=redirect_uri, scope=use_scopes)
+
     @classmethod
     def _jwt_auth(cls, api):
         """JSON Web Token authorization"""
@@ -152,15 +182,24 @@ def destroy(cls):
     def login(cls, auth_type, api):
         cls._init(auth_type, api)
         if auth_type == "code_grant":
-            return cls.get(auth_type, api).authorize(callback=url_for("ds.ds_callback", _external=True))
+            if session.get("pkce_failed", False):
+                return cls.get(auth_type, api).authorize(callback=url_for("ds.ds_callback", _external=True))
+            else:
+                code_verifier = cls.generate_code_verifier()
+                code_challenge = cls.generate_code_challenge(code_verifier)
+                session["code_verifier"] = code_verifier
+                return redirect(cls.get_auth_url_with_pkce(code_challenge))
         elif auth_type == "jwt":
             return cls._jwt_auth(api)
 
     @classmethod
     def get_token(cls, auth_type):
         resp = None
         if auth_type == "code_grant":
-            resp = cls.get(auth_type).authorized_response()
+            if session.get("pkce_failed", False):
+                resp = cls.get(auth_type).authorized_response()
+            else:
+                return cls.fetch_token_with_pkce(request.url)
         elif auth_type == "jwt":
             resp = cls.get(auth_type).to_dict()
 
@@ -189,3 +228,44 @@ def get(cls, auth_type, api=API_TYPE["ESIGNATURE"]):
         if not cls.ds_app:
             cls._init(auth_type, api)
         return cls.ds_app
+    
+    @classmethod
+    def generate_code_verifier(cls):
+        # Generate a random 32-byte string and base64-url encode it
+        return secrets.token_urlsafe(32)
+
+    @classmethod
+    def generate_code_challenge(cls, code_verifier):
+        # Hash the code verifier using SHA-256
+        sha256_hash = hashlib.sha256(code_verifier.encode()).digest()
+
+        # Base64 encode the hash and make it URL safe
+        base64_encoded = base64.urlsafe_b64encode(sha256_hash).decode().rstrip('=')
+
+        return base64_encoded
+    
+    @classmethod
+    def get_auth_url_with_pkce(cls, code_challenge):
+        authorize_url = DS_CONFIG["authorization_server"] + "/oauth/auth"
+        auth_url, state = cls.ds_app.authorization_url(
+            authorize_url,
+            code_challenge=code_challenge,
+            code_challenge_method='S256',  # PKCE uses SHA-256 hashing,
+            approval_prompt="auto"
+        )
+
+        return auth_url
+    
+    @classmethod
+    def fetch_token_with_pkce(cls, authorization_response):
+        access_token_url = DS_CONFIG["authorization_server"] + "/oauth/token"
+        token = cls.get("code_grant", session.get("api")).fetch_token(
+            access_token_url,
+            authorization_response=authorization_response,
+            client_id=DS_CONFIG["ds_client_id"],
+            client_secret=DS_CONFIG["ds_client_secret"],
+            code_verifier=session["code_verifier"],
+            code_challenge_method="S256"
+        )
+
+        return token

app/docusign/views.py

@@ -86,7 +86,14 @@ def ds_callback():
 
     # Save the redirect eg if present
     redirect_url = session.pop("eg", None)
-    resp = DSClient.get_token(session["auth_type"])
+    try:
+        resp = DSClient.get_token(session["auth_type"])
+    except Exception as err:
+        if session.get("pkce_failed", False):
+            raise err
+
+        session["pkce_failed"] = True
+        return redirect(url_for("ds.ds_login"))
 
     # app.logger.info("Authenticated with DocuSign.")
     session["ds_access_token"] = resp["access_token"]

run.py

@@ -1,21 +1,23 @@
-#!flask/bin/python
-from app import app
-from flask_session import Session
-import os
-import sys
-
-host = "0.0.0.0" if "--docker" in sys.argv else "localhost"
-port = int(os.environ.get("PORT", 3000))
-
-if os.environ.get("DEBUG", False) == "True":
-    app.config["DEBUG"] = True
-    app.config['SESSION_TYPE'] = 'filesystem'
-    sess = Session()
-    sess.init_app(app)
-    app.run(host=host, port=port, debug=True)
-else:
-    app.config['SESSION_TYPE'] = 'filesystem'
-    sess = Session()
-    sess.init_app(app)
-    app.run(host=host, port=port, extra_files="api_type.py")
-
+#!flask/bin/python
+from app import app
+from flask_session import Session
+import os
+import sys
+
+host = "0.0.0.0" if "--docker" in sys.argv else "localhost"
+port = int(os.environ.get("PORT", 3000))
+
+os.environ['OAUTHLIB_INSECURE_TRANSPORT'] = '1'
+
+if os.environ.get("DEBUG", False) == "True":
+    app.config["DEBUG"] = True
+    app.config['SESSION_TYPE'] = 'filesystem'
+    sess = Session()
+    sess.init_app(app)
+    app.run(host=host, port=port, debug=True)
+else:
+    app.config['SESSION_TYPE'] = 'filesystem'
+    sess = Session()
+    sess.init_app(app)
+    app.run(host=host, port=port, extra_files="api_type.py")
+