From d198b1abe6b96bbdd84e4c2622d996f770c69edd Mon Sep 17 00:00:00 2001 From: Andrea Grandi Date: Mon, 29 May 2017 17:07:50 +0100 Subject: [PATCH 1/5] Add Django manage command to create a DRF user Token --- .../authtoken/management/__init__.py | 0 .../authtoken/management/commands/__init__.py | 0 .../management/commands/drf_create_token.py | 20 +++++++++++++++++++ tests/test_authtoken.py | 15 ++++++++++++++ 4 files changed, 35 insertions(+) create mode 100644 rest_framework/authtoken/management/__init__.py create mode 100644 rest_framework/authtoken/management/commands/__init__.py create mode 100644 rest_framework/authtoken/management/commands/drf_create_token.py diff --git a/rest_framework/authtoken/management/__init__.py b/rest_framework/authtoken/management/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rest_framework/authtoken/management/commands/__init__.py b/rest_framework/authtoken/management/commands/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rest_framework/authtoken/management/commands/drf_create_token.py b/rest_framework/authtoken/management/commands/drf_create_token.py new file mode 100644 index 000000000..ccd400ed5 --- /dev/null +++ b/rest_framework/authtoken/management/commands/drf_create_token.py @@ -0,0 +1,20 @@ +from django.contrib.auth.models import User +from django.core.management.base import BaseCommand +from rest_framework.authtoken.models import Token + + +class Command(BaseCommand): + help = 'Create DRF Token for a given user' + + def create_user_token(self, username): + user = User.objects.get(username=username) + token = Token.objects.get_or_create(user=user) + return token[0] + + def add_arguments(self, parser): + parser.add_argument('username', type=str, nargs='+') + + def handle(self, *args, **options): + username = options['username'] + token = self.create_user_token(username) + print('Generated token {0} for user {1}'.format(token.key, username)) diff --git a/tests/test_authtoken.py b/tests/test_authtoken.py index 54ac1848d..3b5a618e3 100644 --- a/tests/test_authtoken.py +++ b/tests/test_authtoken.py @@ -4,6 +4,8 @@ from django.contrib.auth.models import User from django.test import TestCase from rest_framework.authtoken.admin import TokenAdmin +from rest_framework.authtoken.management.commands.drf_create_token import \ + Command as AuthTokenCommand from rest_framework.authtoken.models import Token from rest_framework.authtoken.serializers import AuthTokenSerializer from rest_framework.exceptions import ValidationError @@ -33,3 +35,16 @@ class AuthTokenTests(TestCase): self.user.set_password(data['password']) self.user.save() assert AuthTokenSerializer(data=data).is_valid() + + +class AuthTokenCommandTests(TestCase): + + def setUp(self): + self.site = site + self.user = User.objects.create_user(username='test_user') + + def test_command_create_user_token(self): + token = AuthTokenCommand().create_user_token(self.user.username) + assert token is not None + token_saved = Token.objects.first() + assert token.key == token_saved.key From be590d61c09805955d89231cff84fd6fd1ab1ecb Mon Sep 17 00:00:00 2001 From: Andrea Grandi Date: Wed, 31 May 2017 23:01:35 +0100 Subject: [PATCH 2/5] Handle invalid User situation --- .../authtoken/management/commands/drf_create_token.py | 8 +++++++- tests/test_authtoken.py | 4 ++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/rest_framework/authtoken/management/commands/drf_create_token.py b/rest_framework/authtoken/management/commands/drf_create_token.py index ccd400ed5..8419614de 100644 --- a/rest_framework/authtoken/management/commands/drf_create_token.py +++ b/rest_framework/authtoken/management/commands/drf_create_token.py @@ -16,5 +16,11 @@ class Command(BaseCommand): def handle(self, *args, **options): username = options['username'] - token = self.create_user_token(username) + + try: + token = self.create_user_token(username) + except User.DoesNotExist: + print('Cannot create the Token: user {0} does not exist'.format( + username + )) print('Generated token {0} for user {1}'.format(token.key, username)) diff --git a/tests/test_authtoken.py b/tests/test_authtoken.py index 3b5a618e3..407d82944 100644 --- a/tests/test_authtoken.py +++ b/tests/test_authtoken.py @@ -48,3 +48,7 @@ class AuthTokenCommandTests(TestCase): assert token is not None token_saved = Token.objects.first() assert token.key == token_saved.key + + def test_command_create_user_token_invalid_user(self): + with pytest.raises(User.DoesNotExist): + AuthTokenCommand().create_user_token('not_existing_user') From cf196a442490339032ecfeefe4e88581d579e209 Mon Sep 17 00:00:00 2001 From: Andrea Grandi Date: Sat, 3 Jun 2017 10:38:01 +0100 Subject: [PATCH 3/5] Get UserModel from get_user_model and do not infer the natural key --- .../authtoken/management/commands/drf_create_token.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/rest_framework/authtoken/management/commands/drf_create_token.py b/rest_framework/authtoken/management/commands/drf_create_token.py index 8419614de..2dc333a6b 100644 --- a/rest_framework/authtoken/management/commands/drf_create_token.py +++ b/rest_framework/authtoken/management/commands/drf_create_token.py @@ -1,13 +1,16 @@ -from django.contrib.auth.models import User +from django.contrib.auth import get_user_model from django.core.management.base import BaseCommand from rest_framework.authtoken.models import Token +UserModel = get_user_model() + + class Command(BaseCommand): help = 'Create DRF Token for a given user' def create_user_token(self, username): - user = User.objects.get(username=username) + user = UserModel._default_manager.get_by_natural_key(username) token = Token.objects.get_or_create(user=user) return token[0] @@ -19,7 +22,7 @@ class Command(BaseCommand): try: token = self.create_user_token(username) - except User.DoesNotExist: + except UserModel.DoesNotExist: print('Cannot create the Token: user {0} does not exist'.format( username )) From 34c38e0cfe5e880e678704c4d473f082787fca64 Mon Sep 17 00:00:00 2001 From: Andrea Grandi Date: Sat, 3 Jun 2017 11:06:14 +0100 Subject: [PATCH 4/5] Use self.sdtout and CommandError to print output --- .../management/commands/drf_create_token.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/rest_framework/authtoken/management/commands/drf_create_token.py b/rest_framework/authtoken/management/commands/drf_create_token.py index 2dc333a6b..8fe47c44a 100644 --- a/rest_framework/authtoken/management/commands/drf_create_token.py +++ b/rest_framework/authtoken/management/commands/drf_create_token.py @@ -1,5 +1,5 @@ from django.contrib.auth import get_user_model -from django.core.management.base import BaseCommand +from django.core.management.base import BaseCommand, CommandError from rest_framework.authtoken.models import Token @@ -23,7 +23,9 @@ class Command(BaseCommand): try: token = self.create_user_token(username) except UserModel.DoesNotExist: - print('Cannot create the Token: user {0} does not exist'.format( - username - )) - print('Generated token {0} for user {1}'.format(token.key, username)) + raise CommandError( + 'Cannot create the Token: user {0} does not exist'.format( + username) + ) + self.stdout.write( + 'Generated token {0} for user {1}'.format(token.key, username)) From d2459710cae9e17b785e1d67b11d7bc9aa39678a Mon Sep 17 00:00:00 2001 From: Andrea Grandi Date: Sat, 3 Jun 2017 11:58:01 +0100 Subject: [PATCH 5/5] Implement option to reset User token --- .../management/commands/drf_create_token.py | 18 +++++++++++++++-- tests/test_authtoken.py | 20 +++++++++++++++++-- 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/rest_framework/authtoken/management/commands/drf_create_token.py b/rest_framework/authtoken/management/commands/drf_create_token.py index 8fe47c44a..417bdd780 100644 --- a/rest_framework/authtoken/management/commands/drf_create_token.py +++ b/rest_framework/authtoken/management/commands/drf_create_token.py @@ -9,19 +9,33 @@ UserModel = get_user_model() class Command(BaseCommand): help = 'Create DRF Token for a given user' - def create_user_token(self, username): + def create_user_token(self, username, reset_token): user = UserModel._default_manager.get_by_natural_key(username) + + if reset_token: + Token.objects.filter(user=user).delete() + token = Token.objects.get_or_create(user=user) return token[0] def add_arguments(self, parser): parser.add_argument('username', type=str, nargs='+') + parser.add_argument( + '-r', + '--reset', + action='store_true', + dest='reset_token', + default=False, + help='Reset existing User token and create a new one', + ) + def handle(self, *args, **options): username = options['username'] + reset_token = options['reset_token'] try: - token = self.create_user_token(username) + token = self.create_user_token(username, reset_token) except UserModel.DoesNotExist: raise CommandError( 'Cannot create the Token: user {0} does not exist'.format( diff --git a/tests/test_authtoken.py b/tests/test_authtoken.py index 407d82944..6374d7141 100644 --- a/tests/test_authtoken.py +++ b/tests/test_authtoken.py @@ -44,11 +44,27 @@ class AuthTokenCommandTests(TestCase): self.user = User.objects.create_user(username='test_user') def test_command_create_user_token(self): - token = AuthTokenCommand().create_user_token(self.user.username) + token = AuthTokenCommand().create_user_token(self.user.username, False) assert token is not None token_saved = Token.objects.first() assert token.key == token_saved.key def test_command_create_user_token_invalid_user(self): with pytest.raises(User.DoesNotExist): - AuthTokenCommand().create_user_token('not_existing_user') + AuthTokenCommand().create_user_token('not_existing_user', False) + + def test_command_reset_user_token(self): + AuthTokenCommand().create_user_token(self.user.username, False) + first_token_key = Token.objects.first().key + AuthTokenCommand().create_user_token(self.user.username, True) + second_token_key = Token.objects.first().key + + assert first_token_key != second_token_key + + def test_command_do_not_reset_user_token(self): + AuthTokenCommand().create_user_token(self.user.username, False) + first_token_key = Token.objects.first().key + AuthTokenCommand().create_user_token(self.user.username, False) + second_token_key = Token.objects.first().key + + assert first_token_key == second_token_key