Implemented configurable AccessToken and RefreshToken

Author: dirkmoors

oauth2_provider/admin.py

@@ -1,12 +1,14 @@
 from django.contrib import admin
 
-from .models import Grant, AccessToken, RefreshToken, get_application_model
+from .models import Grant, get_access_token_model, get_refresh_token_model, get_application_model
 
 
 class RawIDAdmin(admin.ModelAdmin):
     raw_id_fields = ('user',)
 
 Application = get_application_model()
+AccessToken = get_access_token_model()
+RefreshToken = get_refresh_token_model()
 
 admin.site.register(Application, RawIDAdmin)
 admin.site.register(Grant, RawIDAdmin)

oauth2_provider/models.py

@@ -165,7 +165,7 @@ def __str__(self):
 
 
 @python_2_unicode_compatible
-class AccessToken(models.Model):
+class AbstractAccessToken(models.Model):
     """
     An AccessToken instance represents the actual access token to
     access user's resources, as in :rfc:`5`.
@@ -217,8 +217,15 @@ def __str__(self):
         return self.token
 
 
+class AccessToken(AbstractAccessToken):
+    pass
+
+# Add swappable like this to not break django 1.4 compatibility
+AccessToken._meta.swappable = 'OAUTH2_PROVIDER_ACCESS_TOKEN_MODEL'
+
+
 @python_2_unicode_compatible
-class RefreshToken(models.Model):
+class AbstractRefreshToken(models.Model):
     """
     A RefreshToken instance represents a token that can be swapped for a new
     access token when it expires.
@@ -234,13 +241,20 @@ class RefreshToken(models.Model):
     user = models.ForeignKey(AUTH_USER_MODEL)
     token = models.CharField(max_length=255, db_index=True)
     application = models.ForeignKey(oauth2_settings.APPLICATION_MODEL)
-    access_token = models.OneToOneField(AccessToken,
+    access_token = models.OneToOneField(oauth2_settings.ACCESS_TOKEN_MODEL,
                                         related_name='refresh_token')
 
     def __str__(self):
         return self.token
 
 
+class RefreshToken(AbstractRefreshToken):
+    pass
+
+# Add swappable like this to not break django 1.4 compatibility
+RefreshToken._meta.swappable = 'OAUTH2_PROVIDER_REFRESH_TOKEN_MODEL'
+
+
 def get_application_model():
     """ Return the Application model that is active in this project. """
     try:
@@ -253,3 +267,31 @@ def get_application_model():
         e = "APPLICATION_MODEL refers to model {0} that has not been installed"
         raise ImproperlyConfigured(e.format(oauth2_settings.APPLICATION_MODEL))
     return app_model
+
+
+def get_access_token_model():
+    """ Return the AccessToken model that is active in this project. """
+    try:
+        app_label, model_name = oauth2_settings.ACCESS_TOKEN_MODEL.split('.')
+    except ValueError:
+        e = "ACCESS_TOKEN_MODEL must be of the form 'app_label.model_name'"
+        raise ImproperlyConfigured(e)
+    access_token_model = get_model(app_label, model_name)
+    if access_token_model is None:
+        e = "ACCESS_TOKEN_MODEL refers to model {0} that has not been installed"
+        raise ImproperlyConfigured(e.format(oauth2_settings.ACCESS_TOKEN_MODEL))
+    return access_token_model
+
+
+def get_refresh_token_model():
+    """ Return the RefreshToken model that is active in this project. """
+    try:
+        app_label, model_name = oauth2_settings.REFRESH_TOKEN_MODEL.split('.')
+    except ValueError:
+        e = "REFRESH_TOKEN_MODEL must be of the form 'app_label.model_name'"
+        raise ImproperlyConfigured(e)
+    refresh_token_model = get_model(app_label, model_name)
+    if refresh_token_model is None:
+        e = "REFRESH_TOKEN_MODEL refers to model {0} that has not been installed"
+        raise ImproperlyConfigured(e.format(oauth2_settings.REFRESH_TOKEN_MODEL))
+    return refresh_token_model

oauth2_provider/oauth2_validators.py

@@ -11,7 +11,7 @@
 from oauthlib.oauth2 import RequestValidator
 
 from .compat import unquote_plus
-from .models import Grant, AccessToken, RefreshToken, get_application_model, AbstractApplication
+from .models import Grant, get_access_token_model, get_refresh_token_model, get_application_model, AbstractApplication
 from .settings import oauth2_settings
 
 log = logging.getLogger('oauth2_provider')
@@ -212,6 +212,9 @@ def validate_bearer_token(self, token, scopes, request):
         if not token:
             return False
 
+        # Load the AccessToken model
+        AccessToken = get_access_token_model()
+
         try:
             access_token = AccessToken.objects.select_related("application", "user").get(
                 token=token)
@@ -283,6 +286,13 @@ def save_bearer_token(self, token, request, *args, **kwargs):
         Save access and refresh token, If refresh token is issued, remove old refresh tokens as
         in rfc:`6`
         """
+
+        # Load the AccessToken model
+        AccessToken = get_access_token_model()
+
+        # Load the RefreshToken model
+        RefreshToken = get_refresh_token_model()
+
         if request.refresh_token:
             # remove used refresh token
             try:
@@ -325,6 +335,12 @@ def revoke_token(self, token, token_type_hint, request, *args, **kwargs):
         if token_type_hint not in ['access_token', 'refresh_token']:
             token_type_hint = None
 
+        # Load the AccessToken model
+        AccessToken = get_access_token_model()
+
+        # Load the RefreshToken model
+        RefreshToken = get_refresh_token_model()
+
         token_types = {
             'access_token': AccessToken,
             'refresh_token': RefreshToken,
@@ -358,6 +374,9 @@ def validate_refresh_token(self, refresh_token, client, request, *args, **kwargs
         Check refresh_token exists and refers to the right client.
         Also attach User instance to the request object
         """
+        # Load the RefreshToken model
+        RefreshToken = get_refresh_token_model()
+
         try:
             rt = RefreshToken.objects.get(token=refresh_token)
             request.user = rt.user

oauth2_provider/settings.py

@@ -41,6 +41,8 @@
     'AUTHORIZATION_CODE_EXPIRE_SECONDS': 60,
     'ACCESS_TOKEN_EXPIRE_SECONDS': 36000,
     'APPLICATION_MODEL': getattr(settings, 'OAUTH2_PROVIDER_APPLICATION_MODEL', 'oauth2_provider.Application'),
+    'ACCESS_TOKEN_MODEL': getattr(settings, 'OAUTH2_PROVIDER_ACCESS_TOKEN_MODEL', 'oauth2_provider.AccessToken'),
+    'REFRESH_TOKEN_MODEL': getattr(settings, 'OAUTH2_PROVIDER_REFRESH_TOKEN_MODEL', 'oauth2_provider.RefreshToken'),
     'REQUEST_APPROVAL_PROMPT': 'force',
     'ALLOWED_REDIRECT_URI_SCHEMES': ['http', 'https'],
 

oauth2_provider/tests/test_auth_backends.py

@@ -7,13 +7,13 @@
 
 from ..compat import get_user_model
 from ..models import get_application_model
-from ..models import AccessToken
+from ..models import get_access_token_model
 from ..backends import OAuth2Backend
 from ..middleware import OAuth2TokenMiddleware
 
 UserModel = get_user_model()
 ApplicationModel = get_application_model()
-
+AccessToken = get_access_token_model()
 
 class BaseTest(TestCase):
     """

oauth2_provider/tests/test_authorization_code.py

@@ -10,14 +10,15 @@
 from django.utils import timezone
 
 from ..compat import urlparse, parse_qs, urlencode, get_user_model
-from ..models import get_application_model, Grant, AccessToken
+from ..models import get_application_model, Grant, get_access_token_model
 from ..settings import oauth2_settings
 from ..views import ProtectedResourceView
 
 from .test_utils import TestCaseUtils
 
 
 Application = get_application_model()
+AccessToken = get_access_token_model()
 UserModel = get_user_model()
 
 

oauth2_provider/tests/test_client_credential.py

@@ -13,7 +13,7 @@
 
 from oauthlib.oauth2 import BackendApplicationServer
 
-from ..models import get_application_model, AccessToken
+from ..models import get_application_model, get_access_token_model
 from ..oauth2_backends import OAuthLibCore
 from ..oauth2_validators import OAuth2Validator
 from ..settings import oauth2_settings
@@ -24,6 +24,7 @@
 
 
 Application = get_application_model()
+AccessToken = get_access_token_model()
 UserModel = get_user_model()
 
 

oauth2_provider/tests/test_decorators.py

@@ -6,12 +6,13 @@
 
 from ..decorators import protected_resource, rw_protected_resource
 from ..settings import oauth2_settings
-from ..models import get_application_model, AccessToken
+from ..models import get_application_model, get_access_token_model
 from ..compat import get_user_model
 from .test_utils import TestCaseUtils
 
 
 Application = get_application_model()
+AccessToken = get_access_token_model()
 UserModel = get_user_model()
 
 

oauth2_provider/tests/test_models.py

@@ -11,11 +11,13 @@
 from django.core.exceptions import ValidationError
 from django.utils import timezone
 
-from ..models import get_application_model, Grant, AccessToken, RefreshToken
+from ..models import get_application_model, Grant, get_access_token_model, get_refresh_token_model
 from ..compat import get_user_model
 
 
 Application = get_application_model()
+AccessToken = get_access_token_model()
+RefreshToken = get_refresh_token_model()
 UserModel = get_user_model()
 
 

oauth2_provider/tests/test_rest_framework.py

@@ -7,12 +7,13 @@
 
 
 from .test_utils import TestCaseUtils
-from ..models import AccessToken, get_application_model
+from ..models import get_access_token_model, get_application_model
 from ..settings import oauth2_settings
 from ..compat import get_user_model
 
 
 Application = get_application_model()
+AccessToken = get_access_token_model()
 UserModel = get_user_model()