Move custom validator registration onto GrantTypeBase

6 files changed, 105 insertions, 76 deletions

diff --git a/oauthlib/oauth2/rfc6749/grant_types/authorization_code.py b/oauthlib/oauth2/rfc6749/grant_types/authorization_code.py
index 09756a1..1ac40f3 100644
--- a/oauthlib/oauth2/rfc6749/grant_types/authorization_code.py
+++ b/oauthlib/oauth2/rfc6749/grant_types/authorization_code.py
@@ -13,7 +13,6 @@ from oauthlib.uri_validate import is_absolute_uri
 
 from .base import GrantTypeBase
 from .. import errors
-from ..request_validator import RequestValidator
 
 log = logging.getLogger(__name__)
 
@@ -96,31 +95,7 @@ class AuthorizationCodeGrant(GrantTypeBase):
     """
 
     default_response_mode = 'query'
-
-    def __init__(self, request_validator=None, refresh_token=True):
-        self.request_validator = request_validator or RequestValidator()
-        self.refresh_token = refresh_token
-
-        self._authorization_validators = []
-        self._token_validators = []
-        self._code_modifiers = []
-        self._token_modifiers = []
-        self.response_types = ['code']
-
-    def register_response_type(self, response_type):
-        self.response_types.append(response_type)
-
-    def register_authorization_validator(self, validator):
-        self._authorization_validators.append(validator)
-
-    def register_token_validator(self, validator):
-        self._token_validators.append(validator)
-
-    def register_code_modifier(self, modifier):
-        self._code_modifiers.append(modifier)
-
-    def register_token_modifier(self, modifier):
-        self._token_modifiers.append(modifier)
+    response_types = ['code']
 
     def create_authorization_code(self, request):
         """Generates an authorization grant represented as a dictionary."""
@@ -347,6 +322,10 @@ class AuthorizationCodeGrant(GrantTypeBase):
         # Note that the correct parameters to be added are automatically
         # populated through the use of specific exceptions.
 
+        request_info = {}
+        for validator in self._auth_validators_run_before_standard_ones:
+            request_info.update(validator(request))
+
         # REQUIRED.
         if request.response_type is None:
             raise errors.MissingResponseTypeError(request=request)
@@ -367,15 +346,15 @@ class AuthorizationCodeGrant(GrantTypeBase):
         # http://tools.ietf.org/html/rfc6749#section-3.3
         self.validate_scopes(request)
 
-        request_info = {
+        request_info.update({
             'client_id': request.client_id,
             'redirect_uri': request.redirect_uri,
             'response_type': request.response_type,
             'state': request.state,
             'request': request
-        }
+        })
 
-        for validator in self._authorization_validators:
+        for validator in self._auth_validators_run_after_standard_ones:
             request_info.update(validator(request))
 
         return request.scopes, request_info
@@ -385,6 +364,9 @@ class AuthorizationCodeGrant(GrantTypeBase):
         if request.grant_type not in ('authorization_code', 'openid'):
             raise errors.UnsupportedGrantTypeError(request=request)
 
+        for validator in self._token_validators_run_before_standard_ones:
+            validator(request)
+
         if request.code is None:
             raise errors.InvalidRequestError(
                 description='Missing code parameter.', request=request)
@@ -441,6 +423,5 @@ class AuthorizationCodeGrant(GrantTypeBase):
                       request.redirect_uri, request.client_id, request.client)
             raise errors.AccessDeniedError(request=request)
 
-        for validator in self._token_validators:
+        for validator in self._token_validators_run_after_standard_ones:
             validator(request)
-
diff --git a/oauthlib/oauth2/rfc6749/grant_types/base.py b/oauthlib/oauth2/rfc6749/grant_types/base.py
index 36b06eb..5a3888a 100644
--- a/oauthlib/oauth2/rfc6749/grant_types/base.py
+++ b/oauthlib/oauth2/rfc6749/grant_types/base.py
@@ -9,6 +9,7 @@ import logging
 
 from oauthlib.common import add_params_to_uri
 from oauthlib.oauth2.rfc6749 import errors, utils
+from ..request_validator import RequestValidator
 
 log = logging.getLogger(__name__)
 
@@ -17,6 +18,49 @@ class GrantTypeBase(object):
     error_uri = None
     request_validator = None
     default_response_mode = 'fragment'
+    refresh_token = True
+    response_types = ['code']
+
+    def __init__(self, request_validator=None, **kwargs):
+        self.request_validator = request_validator or RequestValidator()
+
+        # Transforms class variables into instance variables:
+        self.response_types = self.response_types
+        self.refresh_token = self.refresh_token
+
+        self._setup_validator_hooks()
+        for kw, val in kwargs.items():
+            setattr(self, kw, val)
+
+    def _setup_validator_hooks(self):
+        self._auth_validators_run_before_standard_ones = []
+        self._auth_validators_run_after_standard_ones = []
+        self._token_validators_run_before_standard_ones = []
+        self._token_validators_run_after_standard_ones = []
+        self._code_modifiers = []
+        self._token_modifiers = []
+
+    def register_response_type(self, response_type):
+        self.response_types.append(response_type)
+
+    def register_authorization_validator(self, validator, after_standard=True):
+        if after_standard:
+            self._auth_validators_run_after_standard_ones.append(validator)
+        else:
+            self._auth_validators_run_before_standard_ones.append(validator)
+
+    def register_token_validator(self, validator, after_standard=True):
+        if after_standard:
+            self._token_validators_run_after_standard_ones.append(validator)
+        else:
+            self._token_validators_run_before_standard_ones.append(validator)
+
+    def register_code_modifier(self, modifier):
+        self._code_modifiers.append(modifier)
+
+    def register_token_modifier(self, modifier):
+        self._token_modifiers.append(modifier)
+
 
     def create_authorization_response(self, request, token_handler):
         raise NotImplementedError('Subclasses must implement this method.')
diff --git a/oauthlib/oauth2/rfc6749/grant_types/client_credentials.py b/oauthlib/oauth2/rfc6749/grant_types/client_credentials.py
index 91c17a6..1306daf 100644
--- a/oauthlib/oauth2/rfc6749/grant_types/client_credentials.py
+++ b/oauthlib/oauth2/rfc6749/grant_types/client_credentials.py
@@ -50,13 +50,6 @@ class ClientCredentialsGrant(GrantTypeBase):
     .. _`Client Credentials Grant`: http://tools.ietf.org/html/rfc6749#section-4.4
     """
 
-    def __init__(self, request_validator=None):
-        self.request_validator = request_validator or RequestValidator()
-        self._token_modifiers = []
-
-    def register_token_modifier(self, modifier):
-        self._token_modifiers.append(modifier)
-
     def create_token_response(self, request, token_handler):
         """Return token or error in JSON format.
 
@@ -92,6 +85,9 @@ class ClientCredentialsGrant(GrantTypeBase):
         return headers, json.dumps(token), 200
 
     def validate_token_request(self, request):
+        for validator in self._token_validators_run_before_standard_ones:
+            validator(request)
+
         if not getattr(request, 'grant_type', None):
             raise errors.InvalidRequestError('Request is missing grant type.',
                                              request=request)
@@ -119,3 +115,6 @@ class ClientCredentialsGrant(GrantTypeBase):
         log.debug('Authorizing access to user %r.', request.user)
         request.client_id = request.client_id or request.client.client_id
         self.validate_scopes(request)
+
+        for validator in self._token_validators_run_after_standard_ones:
+            validator(request)
diff --git a/oauthlib/oauth2/rfc6749/grant_types/implicit.py b/oauthlib/oauth2/rfc6749/grant_types/implicit.py
index 7366b94..72bb94c 100644
--- a/oauthlib/oauth2/rfc6749/grant_types/implicit.py
+++ b/oauthlib/oauth2/rfc6749/grant_types/implicit.py
@@ -5,6 +5,7 @@ oauthlib.oauth2.rfc6749.grant_types
 """
 from __future__ import unicode_literals, absolute_import
 
+from itertools import chain
 import logging
 
 from oauthlib import common
@@ -117,20 +118,8 @@ class ImplicitGrant(GrantTypeBase):
     .. _`Section 10.16`: http://tools.ietf.org/html/rfc6749#section-10.16
     """
 
-    def __init__(self, request_validator=None):
-        self.request_validator = request_validator or RequestValidator()
-        self._authorization_validators = []
-        self._token_modifiers = []
-        self.response_types = ['token']
-
-    def register_response_type(self, response_type):
-        self.response_types.append(response_type)
-
-    def register_authorization_validator(self, validator):
-        self._authorization_validators.append(validator)
-
-    def register_token_modifier(self, modifier):
-        self._token_modifiers.append(modifier)
+    response_types = ['token']
+    grant_allows_refresh_token = False
 
     def create_authorization_response(self, request, token_handler):
         """Create an authorization response.
@@ -328,6 +317,16 @@ class ImplicitGrant(GrantTypeBase):
 
         # Then check for normal errors.
 
+        request_info = {}
+        # For implicit grant, auth_validators and token_validators are
+        # basically equivalent since the token is returned from the
+        # authorization endpoint.
+        for validator in chain(self._token_validators_run_before_standard_ones,
+                               self._auth_validators_run_before_standard_ones):
+            result = validator(request)
+            if result is not None:
+                request_info.update(result)
+
         # If the resource owner denies the access request or if the request
         # fails for reasons other than a missing or invalid redirection URI,
         # the authorization server informs the client by adding the following
@@ -359,15 +358,21 @@ class ImplicitGrant(GrantTypeBase):
         # http://tools.ietf.org/html/rfc6749#section-3.3
         self.validate_scopes(request)
 
-        request_info = {
+        request_info.update({
                 'client_id': request.client_id,
                 'redirect_uri': request.redirect_uri,
                 'response_type': request.response_type,
                 'state': request.state,
                 'request': request,
-        }
-
-        for validator in self._authorization_validators:
-            request_info.update(validator(request))
+        })
+
+        # For implicit grant, auth_validators and token_validators are
+        # basically equivalent since the token is returned from the
+        # authorization endpoint.
+        for validator in chain(self._auth_validators_run_after_standard_ones,
+                               self._token_validators_run_after_standard_ones):
+            result = validator(request)
+            if result is not None:
+                request_info.update(result)
 
         return request.scopes, request_info
diff --git a/oauthlib/oauth2/rfc6749/grant_types/refresh_token.py b/oauthlib/oauth2/rfc6749/grant_types/refresh_token.py
index cb26880..def630c 100644
--- a/oauthlib/oauth2/rfc6749/grant_types/refresh_token.py
+++ b/oauthlib/oauth2/rfc6749/grant_types/refresh_token.py
@@ -22,13 +22,13 @@ class RefreshTokenGrant(GrantTypeBase):
     .. _`Refresh token grant`: http://tools.ietf.org/html/rfc6749#section-6
     """
 
-    def __init__(self, request_validator=None, issue_new_refresh_tokens=True):
-        self.request_validator = request_validator or RequestValidator()
-        self.issue_new_refresh_tokens = issue_new_refresh_tokens
-        self._token_modifiers = []
-
-    def register_token_modifier(self, modifier):
-        self._token_modifiers.append(modifier)
+    def __init__(self, request_validator=None,
+                 issue_new_refresh_tokens=True,
+                 **kwargs):
+        super(RefreshTokenGrant, self).__init__(
+            request_validator,
+            issue_new_refresh_tokens=issue_new_refresh_tokens,
+            **kwargs)
 
     def create_token_response(self, request, token_handler):
         """Create a new access token from a refresh_token.
@@ -76,6 +76,9 @@ class RefreshTokenGrant(GrantTypeBase):
         if request.grant_type != 'refresh_token':
             raise errors.UnsupportedGrantTypeError(request=request)
 
+        for validator in self._token_validators_run_before_standard_ones:
+            validator(request)
+
         if request.refresh_token is None:
             raise errors.InvalidRequestError(
                 description='Missing refresh token parameter.',
@@ -123,3 +126,6 @@ class RefreshTokenGrant(GrantTypeBase):
                 raise errors.InvalidScopeError(request=request)
         else:
             request.scopes = original_scopes
+
+        for validator in self._token_validators_run_after_standard_ones:
+            validator(request)
diff --git a/oauthlib/oauth2/rfc6749/grant_types/resource_owner_password_credentials.py b/oauthlib/oauth2/rfc6749/grant_types/resource_owner_password_credentials.py
index 0f4d65e..2ef6e16 100644
--- a/oauthlib/oauth2/rfc6749/grant_types/resource_owner_password_credentials.py
+++ b/oauthlib/oauth2/rfc6749/grant_types/resource_owner_password_credentials.py
@@ -70,18 +70,6 @@ class ResourceOwnerPasswordCredentialsGrant(GrantTypeBase):
     .. _`Resource Owner Password Credentials Grant`: http://tools.ietf.org/html/rfc6749#section-4.3
     """
 
-    def __init__(self, request_validator=None, refresh_token=True):
-        """
-        If the refresh_token keyword argument is False, do not return
-        a refresh token in the response.
-        """
-        self.request_validator = request_validator or RequestValidator()
-        self.refresh_token = refresh_token
-        self._token_modifiers = []
-
-    def register_token_modifier(self, modifier):
-        self._token_modifiers.append(modifier)
-
     def create_token_response(self, request, token_handler):
         """Return token or error in json format.
 
@@ -168,6 +156,9 @@ class ResourceOwnerPasswordCredentialsGrant(GrantTypeBase):
         .. _`Section 3.3`: http://tools.ietf.org/html/rfc6749#section-3.3
         .. _`Section 3.2.1`: http://tools.ietf.org/html/rfc6749#section-3.2.1
         """
+        for validator in self._token_validators_run_before_standard_ones:
+            validator(request)
+
         for param in ('grant_type', 'username', 'password'):
             if not getattr(request, param, None):
                 raise errors.InvalidRequestError(
@@ -201,3 +192,6 @@ class ResourceOwnerPasswordCredentialsGrant(GrantTypeBase):
         if request.client:
             request.client_id = request.client_id or request.client.client_id
         self.validate_scopes(request)
+
+        for validator in self._token_validators_run_after_standard_ones:
+            validator(request)