Support for extra token validators and code modifiers.

1 files changed, 15 insertions, 1 deletions

diff --git a/oauthlib/oauth2/rfc6749/grant_types/authorization_code.py b/oauthlib/oauth2/rfc6749/grant_types/authorization_code.py
index f8fe45d..7373f15 100644
--- a/oauthlib/oauth2/rfc6749/grant_types/authorization_code.py
+++ b/oauthlib/oauth2/rfc6749/grant_types/authorization_code.py
@@ -100,6 +100,8 @@ class AuthorizationCodeGrant(GrantTypeBase):
         self.refresh_token = refresh_token
         # NEW-FOR-OPENID
         self._authorization_validators = []
+        self._token_validators = []
+        self._code_modifiers = []
         self._token_modifiers = []
         self.response_types = ['code']
 
@@ -109,6 +111,12 @@ class AuthorizationCodeGrant(GrantTypeBase):
     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)
 
@@ -227,6 +235,8 @@ class AuthorizationCodeGrant(GrantTypeBase):
             return {'Location': common.add_params_to_uri(request.redirect_uri, e.twotuples)}, None, 302
 
         grant = self.create_authorization_code(request)
+        for modifier in self._code_modifiers:
+            grant = modifier(grant, token_handler, request)
         log.debug('Saving grant %r for %r.', grant, request)
         self.request_validator.save_authorization_code(
             request.client_id, grant, request)
@@ -256,7 +266,7 @@ class AuthorizationCodeGrant(GrantTypeBase):
         token = token_handler.create_token(request, refresh_token=self.refresh_token)
         # NEW-FOR-OPENID
         for modifier in self._token_modifiers:
-            token = modifier(token)
+            token = modifier(token, token_handler, request)
         self.request_validator.save_token(token)
         # END-NEW-FOR-OPENID
         self.request_validator.invalidate_authorization_code(
@@ -426,3 +436,7 @@ class AuthorizationCodeGrant(GrantTypeBase):
             log.debug('Redirect_uri (%r) invalid for client %r (%r).',
                       request.redirect_uri, request.client_id, request.client)
             raise errors.AccessDeniedError(request=request)
+
+        for validator in self._token_validators:
+            validator(request)
+