PKCE (#786)

* Added pkce on client side for authorization grant flow. Test cases added

* added new args before kwargs

* updating docstrings with clarification on PKCE params

* adding additional clarification on PKCE parameters

* adding initial function to create code_verifier and tests

* using re.compile for code_verifier allowed characters

* adding initial function to create code_challenge with tests

* replacing appropriate chars for base64 URL

Co-authored-by: Aman Singh Solanki <amans330@gmail.com>

6 files changed, 211 insertions, 6 deletions

diff --git a/oauthlib/oauth2/rfc6749/clients/base.py b/oauthlib/oauth2/rfc6749/clients/base.py
index 35a3fd5..bb4c133 100644
--- a/oauthlib/oauth2/rfc6749/clients/base.py
+++ b/oauthlib/oauth2/rfc6749/clients/base.py
@@ -8,6 +8,10 @@ for consuming OAuth 2.0 RFC6749.
 """
 import time
 import warnings
+import secrets
+import re
+import hashlib
+import base64
 
 from oauthlib.common import generate_token
 from oauthlib.oauth2.rfc6749 import tokens
@@ -61,6 +65,9 @@ class Client:
                  state=None,
                  redirect_url=None,
                  state_generator=generate_token,
+                 code_verifier=None,
+                 code_challenge=None,
+                 code_challenge_method=None,
                  **kwargs):
         """Initialize a client with commonly used attributes.
 
@@ -99,6 +106,15 @@ class Client:
 
         :param state_generator: A no argument state generation callable. Defaults
         to :py:meth:`oauthlib.common.generate_token`.
+
+        :param code_verifier: PKCE parameter. A cryptographically random string that is used to correlate the
+        authorization request to the token request.
+
+        :param code_challenge: PKCE parameter. A challenge derived from the code verifier that is sent in the
+        authorization request, to be verified against later.
+
+        :param code_challenge_method: PKCE parameter. A method that was used to derive code challenge.
+        Defaults to "plain" if not present in the request.
         """
 
         self.client_id = client_id
@@ -113,6 +129,9 @@ class Client:
         self.state_generator = state_generator
         self.state = state
         self.redirect_url = redirect_url
+        self.code_verifier = code_verifier
+        self.code_challenge = code_challenge
+        self.code_challenge_method = code_challenge_method
         self.code = None
         self.expires_in = None
         self._expires_at = None
@@ -471,6 +490,91 @@ class Client:
             raise ValueError("Invalid token placement.")
         return uri, headers, body
 
+    def create_code_verifier(self, length):
+        """Create PKCE **code_verifier** used in computing **code_challenge**. 
+
+           :param length: REQUIRED. The length of the code_verifier.
+
+            The client first creates a code verifier, "code_verifier", for each
+            OAuth 2.0 [RFC6749] Authorization Request, in the following manner:
+
+            code_verifier = high-entropy cryptographic random STRING using the
+            unreserved characters [A-Z] / [a-z] / [0-9] / "-" / "." / "_" / "~"
+            from Section 2.3 of [RFC3986], with a minimum length of 43 characters
+            and a maximum length of 128 characters.
+            
+            .. _`Section 4.1`: https://tools.ietf.org/html/rfc7636#section-4.1
+        """
+        code_verifier = None
+
+        if not length >= 43:
+            raise ValueError("Length must be greater than or equal to 43")
+
+        if not length <= 128:
+            raise ValueError("Length must be less than or equal to 128")
+
+        allowed_characters = re.compile('^[A-Zaa-z0-9-._~]')
+        code_verifier = secrets.token_urlsafe(length)
+
+        if not re.search(allowed_characters, code_verifier):
+            raise ValueError("code_verifier contains invalid characters")
+
+        self.code_verifier = code_verifier
+
+        return code_verifier
+
+    def create_code_challenge(self, code_verifier, code_challenge_method=None):
+        """Create PKCE **code_challenge** derived from the  **code_verifier**.
+
+           :param code_verifier: REQUIRED. The **code_verifier** generated from create_code_verifier().
+           :param code_challenge_method: OPTIONAL. The method used to derive the **code_challenge**. Acceptable
+                values include "S256". DEFAULT is "plain".
+
+
+            The client then creates a code challenge derived from the code
+               verifier by using one of the following transformations on the code
+               verifier:
+
+               plain
+                  code_challenge = code_verifier
+
+               S256
+                  code_challenge = BASE64URL-ENCODE(SHA256(ASCII(code_verifier)))
+
+               If the client is capable of using "S256", it MUST use "S256", as
+               "S256" is Mandatory To Implement (MTI) on the server.  Clients are
+               permitted to use "plain" only if they cannot support "S256" for some
+               technical reason and know via out-of-band configuration that the
+               server supports "plain".
+
+               The plain transformation is for compatibility with existing
+               deployments and for constrained environments that can't use the S256
+               transformation.
+
+            .. _`Section 4.2`: https://tools.ietf.org/html/rfc7636#section-4.2
+        """
+        code_challenge = None
+
+        if code_verifier == None:
+            raise ValueError("Invalid code_verifier")
+
+        if code_challenge_method == None:
+            code_challenge_method = "plain"
+            self.code_challenge_method = code_challenge_method
+            code_challenge = code_verifier
+            self.code_challenge = code_challenge
+
+        if code_challenge_method == "S256":
+            h = hashlib.sha256()
+            h.update(code_verifier.encode(encoding='ascii'))
+            sha256_val = h.digest()
+            code_challenge = bytes.decode(base64.urlsafe_b64encode(sha256_val))
+            # replace '+' with '-', '/' with '_', and remove trailing '='
+            code_challenge = code_challenge.replace("+", "-").replace("/", "_").replace("=", "")
+            self.code_challenge = code_challenge
+
+        return code_challenge
+
     def _add_mac_token(self, uri, http_method='GET', body=None,
                        headers=None, token_placement=AUTH_HEADER, ext=None, **kwargs):
         """Add a MAC token to the request authorization header.
diff --git a/oauthlib/oauth2/rfc6749/clients/web_application.py b/oauthlib/oauth2/rfc6749/clients/web_application.py
index a1f3db1..1d3b2b5 100644
--- a/oauthlib/oauth2/rfc6749/clients/web_application.py
+++ b/oauthlib/oauth2/rfc6749/clients/web_application.py
@@ -41,7 +41,7 @@ class WebApplicationClient(Client):
         self.code = code
 
     def prepare_request_uri(self, uri, redirect_uri=None, scope=None,
-                            state=None, **kwargs):
+                            state=None, code_challenge=None, code_challenge_method='plain', **kwargs):
         """Prepare the authorization code request URI
 
         The client constructs the request URI by adding the following
@@ -62,6 +62,13 @@ class WebApplicationClient(Client):
                         to the client.  The parameter SHOULD be used for preventing
                         cross-site request forgery as described in `Section 10.12`_.
 
+        :param code_challenge: OPTIONAL. PKCE parameter. REQUIRED if PKCE is enforced. 
+                        A challenge derived from the code_verifier that is sent in the 
+                        authorization request, to be verified against later.
+
+        :param code_challenge_method: OPTIONAL. PKCE parameter. A method that was used to derive code challenge.
+                                      Defaults to "plain" if not present in the request.
+
         :param kwargs:  Extra arguments to include in the request URI.
 
         In addition to supplied parameters, OAuthLib will append the ``client_id``
@@ -76,6 +83,10 @@ class WebApplicationClient(Client):
             'https://example.com?client_id=your_id&response_type=code&redirect_uri=https%3A%2F%2Fa.b%2Fcallback'
             >>> client.prepare_request_uri('https://example.com', scope=['profile', 'pictures'])
             'https://example.com?client_id=your_id&response_type=code&scope=profile+pictures'
+            >>> client.prepare_request_uri('https://example.com', code_challenge='kjasBS523KdkAILD2k78NdcJSk2k3KHG6')
+            'https://example.com?client_id=your_id&response_type=code&code_challenge=kjasBS523KdkAILD2k78NdcJSk2k3KHG6'
+            >>> client.prepare_request_uri('https://example.com', code_challenge_method='S256')
+            'https://example.com?client_id=your_id&response_type=code&code_challenge_method=S256'
             >>> client.prepare_request_uri('https://example.com', foo='bar')
             'https://example.com?client_id=your_id&response_type=code&foo=bar'
 
@@ -87,10 +98,11 @@ class WebApplicationClient(Client):
         """
         scope = self.scope if scope is None else scope
         return prepare_grant_uri(uri, self.client_id, 'code',
-                                 redirect_uri=redirect_uri, scope=scope, state=state, **kwargs)
+                                 redirect_uri=redirect_uri, scope=scope, state=state, code_challenge=code_challenge,
+                                 code_challenge_method=code_challenge_method, **kwargs)
 
     def prepare_request_body(self, code=None, redirect_uri=None, body='',
-                             include_client_id=True, **kwargs):
+                             include_client_id=True, code_verifier=None, **kwargs):
         """Prepare the access token request body.
 
         The client makes a request to the token endpoint by adding the
@@ -113,6 +125,9 @@ class WebApplicationClient(Client):
                                   authorization server as described in `Section 3.2.1`_.
         :type include_client_id: Boolean
 
+        :param code_verifier: OPTIONAL. A cryptographically random string that is used to correlate the
+                                        authorization request to the token request.
+
         :param kwargs: Extra parameters to include in the token request.
 
         In addition OAuthLib will add the ``grant_type`` parameter set to
@@ -127,6 +142,8 @@ class WebApplicationClient(Client):
             >>> client = WebApplicationClient('your_id')
             >>> client.prepare_request_body(code='sh35ksdf09sf')
             'grant_type=authorization_code&code=sh35ksdf09sf'
+            >>> client.prepare_request_body(code_verifier='KB46DCKJ873NCGXK5GD682NHDKK34GR')
+            'grant_type=authorization_code&code_verifier=KB46DCKJ873NCGXK5GD682NHDKK34GR'
             >>> client.prepare_request_body(code='sh35ksdf09sf', foo='bar')
             'grant_type=authorization_code&code=sh35ksdf09sf&foo=bar'
 
@@ -154,7 +171,7 @@ class WebApplicationClient(Client):
         kwargs['client_id'] = self.client_id
         kwargs['include_client_id'] = include_client_id
         return prepare_token_request(self.grant_type, code=code, body=body,
-                                     redirect_uri=redirect_uri, **kwargs)
+                                     redirect_uri=redirect_uri, code_verifier=code_verifier, **kwargs)
 
     def parse_request_uri_response(self, uri, state=None):
         """Parse the URI query for code and state.
diff --git a/oauthlib/oauth2/rfc6749/parameters.py b/oauthlib/oauth2/rfc6749/parameters.py
index f07b8bd..44738bb 100644
--- a/oauthlib/oauth2/rfc6749/parameters.py
+++ b/oauthlib/oauth2/rfc6749/parameters.py
@@ -23,7 +23,7 @@ from .utils import is_secure_transport, list_to_scope, scope_to_list
 
 
 def prepare_grant_uri(uri, client_id, response_type, redirect_uri=None,
-                      scope=None, state=None, **kwargs):
+                      scope=None, state=None, code_challenge=None, code_challenge_method='plain', **kwargs):
     """Prepare the authorization grant request URI.
 
     The client constructs the request URI by adding the following
@@ -45,6 +45,11 @@ def prepare_grant_uri(uri, client_id, response_type, redirect_uri=None,
                   back to the client.  The parameter SHOULD be used for
                   preventing cross-site request forgery as described in
                   `Section 10.12`_.
+    :param code_challenge: PKCE paramater. A challenge derived from the 
+                           code_verifier that is sent in the authorization 
+                           request, to be verified against later.
+    :param code_challenge_method: PKCE parameter. A method that was used to derive the 
+                                  code_challenge. Defaults to "plain" if not present in the request.
     :param kwargs: Extra arguments to embed in the grant/authorization URL.
 
     An example of an authorization code grant authorization URL:
@@ -52,6 +57,7 @@ def prepare_grant_uri(uri, client_id, response_type, redirect_uri=None,
     .. code-block:: http
 
         GET /authorize?response_type=code&client_id=s6BhdRkqt3&state=xyz
+            &code_challenge=kjasBS523KdkAILD2k78NdcJSk2k3KHG6&code_challenge_method=S256
             &redirect_uri=https%3A%2F%2Fclient%2Eexample%2Ecom%2Fcb HTTP/1.1
         Host: server.example.com
 
@@ -73,6 +79,9 @@ def prepare_grant_uri(uri, client_id, response_type, redirect_uri=None,
         params.append(('scope', list_to_scope(scope)))
     if state:
         params.append(('state', state))
+    if code_challenge is not None:
+        params.append(('code_challenge', code_challenge))
+        params.append(('code_challenge_method', code_challenge_method))
 
     for k in kwargs:
         if kwargs[k]:
@@ -81,7 +90,7 @@ def prepare_grant_uri(uri, client_id, response_type, redirect_uri=None,
     return add_params_to_uri(uri, params)
 
 
-def prepare_token_request(grant_type, body='', include_client_id=True, **kwargs):
+def prepare_token_request(grant_type, body='', include_client_id=True, code_verifier=None, **kwargs):
     """Prepare the access token request.
 
     The client makes a request to the token endpoint by adding the
@@ -116,6 +125,9 @@ def prepare_token_request(grant_type, body='', include_client_id=True, **kwargs)
                          authorization request as described in
                          `Section 4.1.1`_, and their values MUST be identical. *
 
+    :param code_verifier: PKCE parameter. A cryptographically random string that is used to correlate the
+                          authorization request to the token request.
+
     :param kwargs: Extra arguments to embed in the request body.
 
     Parameters marked with a `*` above are not explicit arguments in the
@@ -142,6 +154,10 @@ def prepare_token_request(grant_type, body='', include_client_id=True, **kwargs)
         if client_id is not None:
             params.append(('client_id', client_id))
 
+    # use code_verifier if code_challenge was passed in the authorization request
+    if code_verifier is not None:
+        params.append(('code_verifier', code_verifier))
+
     # the kwargs iteration below only supports including boolean truth (truthy)
     # values, but some servers may require an empty string for `client_secret`
     client_secret = kwargs.pop('client_secret', None)
diff --git a/tests/oauth2/rfc6749/clients/test_base.py b/tests/oauth2/rfc6749/clients/test_base.py
index 6b4eff0..70a2283 100644
--- a/tests/oauth2/rfc6749/clients/test_base.py
+++ b/tests/oauth2/rfc6749/clients/test_base.py
@@ -325,3 +325,31 @@ class ClientTest(TestCase):
         self.assertEqual(client.access_token, response.get("access_token"))
         self.assertEqual(client.refresh_token, response.get("refresh_token"))
         self.assertEqual(client.token_type, response.get("token_type"))
+
+
+    def test_create_code_verifier_min_length(self):
+        client = Client(self.client_id)
+        length = 43
+        code_verifier = client.create_code_verifier(length=length)
+        self.assertEqual(client.code_verifier, code_verifier)
+
+    def test_create_code_verifier_max_length(self):
+        client = Client(self.client_id)
+        length = 128
+        code_verifier = client.create_code_verifier(length=length)
+        self.assertEqual(client.code_verifier, code_verifier)
+
+    def test_create_code_challenge_plain(self):
+        client = Client(self.client_id)
+        code_verifier = client.create_code_verifier(length=128)
+        code_challenge_plain = client.create_code_challenge(code_verifier=code_verifier)
+
+        # if no code_challenge_method specified, code_challenge = code_verifier
+        self.assertEqual(code_challenge_plain, client.code_verifier)
+        self.assertEqual(client.code_challenge_method, "plain")
+
+    def test_create_code_challenge_s256(self):
+        client = Client(self.client_id)
+        code_verifier = client.create_code_verifier(length=128)
+        code_challenge_s256 = client.create_code_challenge(code_verifier=code_verifier, code_challenge_method='S256')
+        self.assertEqual(code_challenge_s256, client.code_challenge)
diff --git a/tests/oauth2/rfc6749/clients/test_web_application.py b/tests/oauth2/rfc6749/clients/test_web_application.py
index 1f711f4..f6b9449 100644
--- a/tests/oauth2/rfc6749/clients/test_web_application.py
+++ b/tests/oauth2/rfc6749/clients/test_web_application.py
@@ -24,10 +24,15 @@ class WebApplicationClientTest(TestCase):
     uri_id = uri + "&response_type=code&client_id=" + client_id
     uri_redirect = uri_id + "&redirect_uri=http%3A%2F%2Fmy.page.com%2Fcallback"
     redirect_uri = "http://my.page.com/callback"
+    code_verifier = "code_verifier"
     scope = ["/profile"]
     state = "xyz"
+    code_challenge = "code_challenge"
+    code_challenge_method = "S256"
     uri_scope = uri_id + "&scope=%2Fprofile"
     uri_state = uri_id + "&state=" + state
+    uri_code_challenge = uri_id + "&code_challenge=" + code_challenge + "&code_challenge_method=" + code_challenge_method
+    uri_code_challenge_method = uri_id + "&code_challenge=" + code_challenge + "&code_challenge_method=plain"
     kwargs = {
         "some": "providers",
         "require": "extra arguments"
@@ -40,6 +45,7 @@ class WebApplicationClientTest(TestCase):
 
     body_code = "not=empty&grant_type=authorization_code&code={}&client_id={}".format(code, client_id)
     body_redirect = body_code + "&redirect_uri=http%3A%2F%2Fmy.page.com%2Fcallback"
+    bode_code_verifier = body_code + "&code_verifier=code_verifier"
     body_kwargs = body_code + "&some=providers&require=extra+arguments"
 
     response_uri = "https://client.example.com/cb?code=zzzzaaaa&state=xyz"
@@ -80,6 +86,14 @@ class WebApplicationClientTest(TestCase):
         uri = client.prepare_request_uri(self.uri, state=self.state)
         self.assertURLEqual(uri, self.uri_state)
 
+        # with code_challenge and code_challenge_method
+        uri = client.prepare_request_uri(self.uri, code_challenge=self.code_challenge, code_challenge_method=self.code_challenge_method)
+        self.assertURLEqual(uri, self.uri_code_challenge)
+
+        # with no code_challenge_method
+        uri = client.prepare_request_uri(self.uri, code_challenge=self.code_challenge)
+        self.assertURLEqual(uri, self.uri_code_challenge_method)
+
         # With extra parameters through kwargs
         uri = client.prepare_request_uri(self.uri, **self.kwargs)
         self.assertURLEqual(uri, self.uri_kwargs)
@@ -99,6 +113,10 @@ class WebApplicationClientTest(TestCase):
         body = client.prepare_request_body(body=self.body, redirect_uri=self.redirect_uri)
         self.assertFormBodyEqual(body, self.body_redirect)
 
+        # With code verifier
+        body = client.prepare_request_body(body=self.body, code_verifier=self.code_verifier)
+        self.assertFormBodyEqual(body, self.bode_code_verifier)
+
         # With extra parameters
         body = client.prepare_request_body(body=self.body, **self.kwargs)
         self.assertFormBodyEqual(body, self.body_kwargs)
diff --git a/tests/oauth2/rfc6749/test_parameters.py b/tests/oauth2/rfc6749/test_parameters.py
index f9245ec..cd8c9e9 100644
--- a/tests/oauth2/rfc6749/test_parameters.py
+++ b/tests/oauth2/rfc6749/test_parameters.py
@@ -21,12 +21,15 @@ class ParameterTests(TestCase):
     list_scope = ['list', 'of', 'scopes']
 
     auth_grant = {'response_type': 'code'}
+    auth_grant_pkce = {'response_type': 'code', 'code_challenge': "code_challenge",
+                       'code_challenge_method': 'code_challenge_method'}
     auth_grant_list_scope = {}
     auth_implicit = {'response_type': 'token', 'extra': 'extra'}
     auth_implicit_list_scope = {}
 
     def setUp(self):
         self.auth_grant.update(self.auth_base)
+        self.auth_grant_pkce.update(self.auth_base)
         self.auth_implicit.update(self.auth_base)
         self.auth_grant_list_scope.update(self.auth_grant)
         self.auth_grant_list_scope['scope'] = self.list_scope
@@ -37,7 +40,14 @@ class ParameterTests(TestCase):
                      '&client_id=s6BhdRkqt3&redirect_uri=https%3A%2F%2F'
                      'client.example.com%2Fcb&scope={1}&state={2}{3}')
 
+    auth_base_uri_pkce = ('https://server.example.com/authorize?response_type={0}'
+                     '&client_id=s6BhdRkqt3&redirect_uri=https%3A%2F%2F'
+                     'client.example.com%2Fcb&scope={1}&state={2}{3}&code_challenge={4}'
+                     '&code_challenge_method={5}')
+
     auth_grant_uri = auth_base_uri.format('code', 'photos', state, '')
+    auth_grant_uri_pkce = auth_base_uri_pkce.format('code', 'photos', state, '', 'code_challenge',
+                                                    'code_challenge_method')
     auth_grant_uri_list_scope = auth_base_uri.format('code', 'list+of+scopes', state, '')
     auth_implicit_uri = auth_base_uri.format('token', 'photos', state, '&extra=extra')
     auth_implicit_uri_list_scope = auth_base_uri.format('token', 'list+of+scopes', state, '&extra=extra')
@@ -47,11 +57,21 @@ class ParameterTests(TestCase):
         'code': 'SplxlOBeZQQYbYS6WxSbIA',
         'redirect_uri': 'https://client.example.com/cb'
     }
+    grant_body_pkce = {
+        'grant_type': 'authorization_code',
+        'code': 'SplxlOBeZQQYbYS6WxSbIA',
+        'redirect_uri': 'https://client.example.com/cb',
+        'code_verifier': 'code_verifier'
+    }
     grant_body_scope = {'scope': 'photos'}
     grant_body_list_scope = {'scope': list_scope}
     auth_grant_body = ('grant_type=authorization_code&'
                        'code=SplxlOBeZQQYbYS6WxSbIA&'
                        'redirect_uri=https%3A%2F%2Fclient.example.com%2Fcb')
+    auth_grant_body_pkce = ('grant_type=authorization_code&'
+                       'code=SplxlOBeZQQYbYS6WxSbIA&'
+                       'redirect_uri=https%3A%2F%2Fclient.example.com%2Fcb'
+                       '&code_verifier=code_verifier')
     auth_grant_body_scope = auth_grant_body + '&scope=photos'
     auth_grant_body_list_scope = auth_grant_body + '&scope=list+of+scopes'
 
@@ -179,12 +199,14 @@ class ParameterTests(TestCase):
         self.assertURLEqual(prepare_grant_uri(**self.auth_grant_list_scope), self.auth_grant_uri_list_scope)
         self.assertURLEqual(prepare_grant_uri(**self.auth_implicit), self.auth_implicit_uri)
         self.assertURLEqual(prepare_grant_uri(**self.auth_implicit_list_scope), self.auth_implicit_uri_list_scope)
+        self.assertURLEqual(prepare_grant_uri(**self.auth_grant_pkce), self.auth_grant_uri_pkce)
 
     def test_prepare_token_request(self):
         """Verify correct access token request body construction."""
         self.assertFormBodyEqual(prepare_token_request(**self.grant_body), self.auth_grant_body)
         self.assertFormBodyEqual(prepare_token_request(**self.pwd_body), self.password_body)
         self.assertFormBodyEqual(prepare_token_request(**self.cred_grant), self.cred_body)
+        self.assertFormBodyEqual(prepare_token_request(**self.grant_body_pkce), self.auth_grant_body_pkce)
 
     def test_grant_response(self):
         """Verify correct parameter parsing and validation for auth code responses."""