Refactor 2FA and prepare for fido support

ractf · Aug 10, 2020 · c57a4d1 · c57a4d1
1 parent e5dee96
commit c57a4d1
Show file tree

Hide file tree

Showing 9 changed files with 198 additions and 60 deletions.
diff --git a/src/authentication/basic_auth.py b/src/authentication/basic_auth.py
@@ -64,12 +64,7 @@ def register_user(self, username, email, password, invite, **kwargs):
 class BasicAuthLoginProvider(LoginProvider):
     name = 'basic_auth'
 
-    #TODO: These fields don't do anything yet, but they will eventually correlate to the kwargs in login_user
-    username = serializers.CharField(max_length=50)
-    password = serializers.CharField(max_length=50)
-    otp = serializers.CharField(max_length=6)
-
-    def login_user(self, username, password, otp, context, **kwargs):
+    def login_user(self, username, password, context, **kwargs):
         user = authenticate(request=context.get('request'),
                             username=username, password=password)
         if not user:
@@ -87,16 +82,6 @@ def login_user(self, username, password, otp, context, **kwargs):
             raise FormattedException(m='login_not_open', d={'reason': 'login_not_open'},
                                      status_code=HTTP_401_UNAUTHORIZED)
 
-        if user.totp_status == TOTPStatus.ENABLED:
-            if not otp or otp == '':
-                login_reject.send(sender=self.__class__, username=username, reason='no_2fa')
-                raise FormattedException(m='2fa_required', d={'reason': '2fa_required'},
-                                         status_code=HTTP_401_UNAUTHORIZED)
-            totp = pyotp.TOTP(user.totp_secret)
-            if not totp.verify(otp, valid_window=1):
-                login_reject.send(sender=self.__class__, username=username, reason='incorrect_2fa')
-                raise FormattedException(m='incorrect_2fa', d={'reason': 'incorrect_2fa'},
-                                         status_code=HTTP_401_UNAUTHORIZED)
         login.send(sender=self.__class__, user=user)
         return user
 

diff --git a/src/authentication/migrations/0004_auto_20200810_1111.py b/src/authentication/migrations/0004_auto_20200810_1111.py
@@ -0,0 +1,44 @@
+# Generated by Django 3.0.5 on 2020-08-10 11:11
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+import pyotp
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('authentication', '0003_passwordresettoken'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='passwordresettoken',
+            name='issued',
+            field=models.DateTimeField(auto_now_add=True),
+        ),
+        migrations.CreateModel(
+            name='TOTPDevice',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('created', models.DateTimeField(auto_now_add=True)),
+                ('last_used', models.DateTimeField(null=True)),
+                ('totp_secret', models.CharField(default=pyotp.random_base32, max_length=16, null=True)),
+                ('verified', models.BooleanField(default=False)),
+                ('user', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, related_name='totp_device', to=settings.AUTH_USER_MODEL)),
+            ],
+        ),
+        migrations.CreateModel(
+            name='BackupCode',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('code', models.CharField(max_length=8)),
+                ('user', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='backup_codes', to=settings.AUTH_USER_MODEL)),
+            ],
+            options={
+                'unique_together': {('user', 'code')},
+            },
+        ),
+    ]
diff --git a/src/authentication/models.py b/src/authentication/models.py
@@ -1,5 +1,6 @@
 from datetime import timedelta
 
+import pyotp
 from django.contrib.auth import get_user_model
 from django.db import models
 from django.utils import timezone
@@ -22,5 +23,33 @@ def one_day():
 class PasswordResetToken(models.Model):
     user = models.ForeignKey(get_user_model(), on_delete=models.CASCADE)
     token = models.CharField(max_length=64)
-    issued = models.DateTimeField(default=timezone.now)
+    issued = models.DateTimeField(auto_now_add=True)
     expires = models.DateTimeField(default=one_day)
+
+
+class BackupCode(models.Model):
+    user = models.ForeignKey(get_user_model(), related_name='backup_codes', on_delete=models.CASCADE)
+    code = models.CharField(max_length=8)
+
+    class Meta:
+        unique_together = [
+            ('user', 'code')
+        ]
+
+    @staticmethod
+    def generate(user):
+        BackupCode.objects.filter(user=user).delete()
+        codes = [BackupCode(user=user, code=pyotp.random_base32(8)) for i in range(10)]
+        BackupCode.objects.bulk_create(codes)
+        return BackupCode.objects.filter(user=user).values_list('code', flat=True)
+
+
+class TOTPDevice(models.Model):
+    user = models.OneToOneField(get_user_model(), related_name='totp_device', on_delete=models.CASCADE)
+    created = models.DateTimeField(auto_now_add=True)
+    last_used = models.DateTimeField(null=True)
+    totp_secret = models.CharField(null=True, max_length=16, default=pyotp.random_base32)
+    verified = models.BooleanField(default=False)
+
+    def validate_token(self, token):
+        return pyotp.TOTP(self.totp_secret).verify(token, valid_window=1)
diff --git a/src/authentication/serializers.py b/src/authentication/serializers.py
@@ -12,12 +12,21 @@
 class LoginSerializer(serializers.Serializer):
     username = serializers.CharField()
     password = serializers.CharField(trim_whitespace=False)
-    otp = serializers.CharField(max_length=6, allow_null=True, allow_blank=True)
 
     def validate(self, data):
         user = providers.get_provider('login').login_user(**data, context=self.context)
-        if user is not None:
-            data['user'] = user
+        data['user'] = user
+        return data
+
+
+class LoginTwoFactorSerializer(serializers.Serializer):
+    username = serializers.CharField()
+    password = serializers.CharField(trim_whitespace=False)
+    tfa = serializers.CharField(max_length=255, allow_null=True, allow_blank=True)
+
+    def validate(self, data):
+        user = providers.get_provider('login').login_user(**data, context=self.context)
+        data['user'] = user
         return data
 
 

diff --git a/src/authentication/urls.py b/src/authentication/urls.py
@@ -12,12 +12,14 @@
     path('add_2fa/', views.AddTwoFactorView.as_view(), name='add-2fa'),
     path('verify_2fa/', views.VerifyTwoFactorView.as_view(), name='verify-2fa'),
     path('remove_2fa/', views.VerifyTwoFactorView.as_view(), name='remove-2fa'),
+    path('login_2fa/', views.LoginTwoFactorView.as_view(), name='login-2fa'),
     path('logout/', views.LogoutView.as_view(), name='logout'),
     path('request_password_reset/', views.RequestPasswordResetView.as_view(), name='request-password-reset'),
     path('password_reset/', views.DoPasswordResetView.as_view(), name='do-password-reset'),
     path('verify_email/', views.VerifyEmailView.as_view(), name='verify-email'),
     path('resend_email/', views.ResendEmailView.as_view(), name='resend-email'),
     path('change_password/', views.ChangePasswordView.as_view(), name='change-password'),
     path('generate_invites/', views.GenerateInvitesView.as_view(), name='generate-invites'),
-    path('invites/', include(router.urls), name='invites')
+    path('invites/', include(router.urls), name='invites'),
+    path('regenerate_backup_codes', views.RegenerateBackupCodesView.as_view(), name='regenerate-backup-codes')
 ]
diff --git a/src/authentication/views.py b/src/authentication/views.py
@@ -1,43 +1,39 @@
 import random
-import string
 import secrets
+import string
 
-import pyotp
 from django.contrib.auth import get_user_model
 from django.core.validators import EmailValidator
 from django.db import transaction
 from django.utils.decorators import method_decorator
 from django.views.decorators.debug import sensitive_post_parameters
 from django_filters.rest_framework import DjangoFilterBackend
 from rest_framework import permissions
-from rest_framework.authtoken.views import ObtainAuthToken
 from rest_framework.generics import CreateAPIView, GenericAPIView, get_object_or_404
 from rest_framework.status import HTTP_400_BAD_REQUEST, HTTP_401_UNAUTHORIZED
 from rest_framework.views import APIView
 
 from authentication import serializers
+from authentication.models import InviteCode, PasswordResetToken, TOTPDevice, BackupCode
 from authentication.permissions import HasTwoFactor, VerifyingTwoFactor
 from authentication.serializers import RegistrationSerializer, EmailVerificationSerializer, ChangePasswordSerializer, \
     GenerateInvitesSerializer, InviteCodeSerializer, EmailSerializer
-from backend import renderers
 from backend.mail import send_email
 from backend.response import FormattedResponse
 from backend.signals import logout, add_2fa, verify_2fa, password_reset_start, password_reset_start_reject, \
     email_verified, change_password, password_reset, remove_2fa
 from backend.viewsets import AdminListModelViewSet
 from member.models import TOTPStatus
-from team.models import Team
-from authentication.models import InviteCode, PasswordResetToken
 from plugins import providers
+from team.models import Team
 
-hide_password = method_decorator(sensitive_post_parameters("password",))
+hide_password = method_decorator(sensitive_post_parameters("password", ))
 
 
 class LoginView(APIView):
     permission_classes = (~permissions.IsAuthenticated,)
     serializer_class = serializers.LoginSerializer
     throttle_scope = "login"
-    renderer_classes = (renderers.RACTFJSONRenderer,)
 
     @hide_password
     def dispatch(self, *args, **kwargs):
@@ -50,6 +46,9 @@ def post(self, request, *args, **kwargs):
         if user is None:
             return FormattedResponse(status=HTTP_401_UNAUTHORIZED, d={'reason': 'login_failed'}, m='login_failed')
 
+        if user.has_2fa():
+            return FormattedResponse(status=HTTP_401_UNAUTHORIZED, d={'reason': '2fa_required'}, m='2fa_required')
+
         token = providers.get_provider('token').issue_token(user)
         return FormattedResponse({'token': token})
 
@@ -79,34 +78,32 @@ class AddTwoFactorView(APIView):
     throttle_scope = "2fa"
 
     def post(self, request):
-        totp_secret = pyotp.random_base32()
-        request.user.totp_secret = totp_secret
-        request.user.totp_status = TOTPStatus.VERIFYING
-        request.user.save()
+        totp_device = TOTPDevice(user=request.user)
+        totp_device.save()
         add_2fa.send(sender=self.__class__, user=request.user)
-        return FormattedResponse({"totp_secret": totp_secret})
+        return FormattedResponse({"totp_secret": totp_device.totp_secret})
 
 
 class VerifyTwoFactorView(APIView):
     permission_classes = (permissions.IsAuthenticated & VerifyingTwoFactor,)
     throttle_scope = "2fa"
 
     def post(self, request):
-        totp = pyotp.TOTP(request.user.totp_secret)
-        valid = totp.verify(request.data["otp"], valid_window=1)
-        if valid:
-            request.user.totp_status = TOTPStatus.ENABLED
-            request.user.save()
+        if request.user.totp_device is not None and request.user.totp_device.validate_token(request.data["otp"]):
+            request.user.totp_device.verified = True
+            request.user.totp_device.save()
+            backup_codes = BackupCode.generate(request.user)
             verify_2fa.send(sender=self.__class__, user=request.user)
-        return FormattedResponse({"valid": valid})
+            return FormattedResponse({"valid": True, "backup_codes": backup_codes})
+        return FormattedResponse({"valid": False})
 
 
 class RemoveTwoFactorView(APIView):
     permission_classes = (permissions.IsAuthenticated & HasTwoFactor,)
     throttle_scope = "2fa"
 
     def post(self, request):
-        request.user.totp_status = TOTPStatus.DISABLED
+        request.user.totp_device.delete()
         request.user.save()
         remove_2fa.send(sender=self.__class__, user=request.user)
         send_email(
@@ -117,6 +114,54 @@ def post(self, request):
         return FormattedResponse()
 
 
+class LoginTwoFactorView(APIView):
+    permission_classes = (~permissions.IsAuthenticated,)
+    serializer_class = serializers.LoginTwoFactorSerializer
+    throttle_scope = "login"
+
+    @hide_password
+    def dispatch(self, *args, **kwargs):
+        return super(LoginTwoFactorView, self).dispatch(*args, **kwargs)
+
+    def issue_token(self, user):
+        token = providers.get_provider('token').issue_token(user)
+        return FormattedResponse({'token': token})
+
+    def post(self, request, *args, **kwargs):
+        serializer = self.serializer_class(data=request.data, context={'request': request})
+        serializer.is_valid(raise_exception=True)
+        user = serializer.validated_data['user']
+        if user is None:
+            return FormattedResponse(status=HTTP_401_UNAUTHORIZED, d={'reason': 'login_failed'}, m='login_failed')
+
+        if user.totp_status != TOTPStatus.ENABLED:
+            return FormattedResponse(status=HTTP_401_UNAUTHORIZED, d={'reason': '2fa_not_enabled'}, m='2fa_not_enabled')
+
+        token = serializer.data['tfa']
+
+        if len(token) == 6:
+            for device in user.totp_devices:
+                if device.validate_token(token):
+                    return self.issue_token(user)
+        elif len(token) == 8:
+            for code in user.backup_codes:
+                if token == code.code:
+                    code.delete()
+                    return self.issue_token(user)
+
+        return self.issue_token(user)
+
+
+class RegenerateBackupCodesView(APIView):
+    permission_classes = (permissions.IsAuthenticated & HasTwoFactor,)
+    serializer_class = serializers.LoginTwoFactorSerializer
+    throttle_scope = "2fa"
+
+    def post(self, request, *args, **kwargs):
+        backup_codes = BackupCode.generate(request.user)
+        return FormattedResponse({"backup_codes": backup_codes})
+
+
 class RequestPasswordResetView(APIView):
     permission_classes = (~permissions.IsAuthenticated,)
     throttle_scope = "request_password_reset"
@@ -265,8 +310,8 @@ def post(self, request):
         if serializer.validated_data["auto_team"]:
             team = get_object_or_404(Team, id=serializer.validated_data["auto_team"])
         with transaction.atomic():
-            for i in range(active_codes, serializer.validated_data["amount"]+active_codes):
-                code = f"{''.join([random.choice(string.ascii_letters+string.digits) for _ in range(8)])}{hex(i)[2:]}"
+            for i in range(active_codes, serializer.validated_data["amount"] + active_codes):
+                code = f"{''.join([random.choice(string.ascii_letters + string.digits) for _ in range(8)])}{hex(i)[2:]}"
                 codes.append(code)
                 invite = InviteCode(code=code, max_uses=serializer.validated_data["max_uses"])
                 if serializer.validated_data["auto_team"]:

diff --git a/src/member/migrations/0004_auto_20200810_1111.py b/src/member/migrations/0004_auto_20200810_1111.py
@@ -0,0 +1,26 @@
+# Generated by Django 3.0.5 on 2020-08-10 11:11
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('member', '0003_auto_20200808_1649'),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name='member',
+            name='totp_secret',
+        ),
+        migrations.RemoveField(
+            model_name='member',
+            name='totp_status',
+        ),
+        migrations.AddField(
+            model_name='member',
+            name='state_actor',
+            field=models.BooleanField(default=False),
+        ),
+    ]
diff --git a/src/member/models.py b/src/member/models.py
@@ -1,3 +1,4 @@
+import random
 import secrets
 import time
 from enum import IntEnum
@@ -32,14 +33,10 @@ class Member(AbstractUser):
             "Required. 36 characters or fewer. Letters, digits and @/./+/-/_ only."
         ),
         validators=[username_validator],
-        error_messages={"unique": _("A user with that username already exists."),},
+        error_messages={"unique": _("A user with that username already exists.")},
     )
     email = models.EmailField(_("email address"), blank=True, unique=True)
-    totp_secret = models.CharField(null=True, max_length=16)
-    totp_status = models.IntegerField(
-        choices=[(status, status.value) for status in TOTPStatus],
-        default=TOTPStatus.DISABLED,
-    )
+    state_actor = models.BooleanField(default=False)
     is_visible = models.BooleanField(default=False)
     bio = models.TextField(blank=True, max_length=400)
     discord = models.CharField(blank=True, max_length=36)
@@ -70,11 +67,11 @@ def issue_token(self):
         token, created = Token.objects.get_or_create(user=self)
         return token.key
 
-    def is_2fa_enabled(self):
-        return self.totp_status == TOTPStatus.ENABLED
+    def has_2fa(self):
+        return self.totp_device is not None and self.totp_device.verified
 
     def should_deny_admin(self):
-        return self.totp_status != TOTPStatus.ENABLED and config.get(
+        return not self.has_2fa() and config.get(
             "enable_force_admin_2fa"
         )