2fa improvements (baserow#4171)

Author: stribny

backend/src/baserow/api/two_factor_auth/errors.py

@@ -3,6 +3,7 @@
     HTTP_401_UNAUTHORIZED,
     HTTP_403_FORBIDDEN,
     HTTP_404_NOT_FOUND,
+    HTTP_429_TOO_MANY_REQUESTS,
 )
 
 ERROR_TWO_FACTOR_AUTH_TYPE_DOES_NOT_EXIST = (
@@ -17,6 +18,12 @@
     "Two-factor authentication verification failed.",
 )
 
+ERROR_RATE_LIMIT_EXCEEDED = (
+    "ERROR_RATE_LIMIT_EXCEEDED",
+    HTTP_429_TOO_MANY_REQUESTS,
+    "Rate limit exceeded.",
+)
+
 ERROR_WRONG_PASSWORD = (
     "ERROR_WRONG_PASSWORD",
     HTTP_403_FORBIDDEN,

backend/src/baserow/api/two_factor_auth/views.py

@@ -13,6 +13,7 @@
 )
 from baserow.api.schemas import get_error_schema
 from baserow.api.two_factor_auth.errors import (
+    ERROR_RATE_LIMIT_EXCEEDED,
     ERROR_TWO_FACTOR_AUTH_ALREADY_CONFIGURED,
     ERROR_TWO_FACTOR_AUTH_CANNOT_BE_CONFIGURED,
     ERROR_TWO_FACTOR_AUTH_NOT_CONFIGURED,
@@ -27,7 +28,7 @@
     VerifyTOTPSerializer,
 )
 from baserow.api.two_factor_auth.tokens import Require2faToken
-from baserow.api.user.schemas import create_user_response_schema
+from baserow.api.user.schemas import authenticated_user_response_schema
 from baserow.api.user.serializers import log_in_user
 from baserow.api.utils import DiscriminatorCustomFieldsMappingSerializer
 from baserow.core.models import User
@@ -48,6 +49,8 @@
     TOTPAuthProviderType,
     two_factor_auth_type_registry,
 )
+from baserow.throttling import RateLimitExceededException, rate_limit
+from baserow.throttling_types import RateLimit
 
 
 class ConfigureTwoFactorAuthView(APIView):
@@ -181,20 +184,22 @@ class VerifyTOTPAuthView(APIView):
         description=("Verifies TOTP two-factor authentication"),
         request=VerifyTOTPSerializer,
         responses={
-            200: create_user_response_schema,
+            200: authenticated_user_response_schema,
             400: get_error_schema(
                 [
                     "ERROR_REQUEST_BODY_VALIDATION",
                 ]
             ),
             401: get_error_schema(["ERROR_TWO_FACTOR_AUTH_VERIFICATION_FAILED"]),
             404: get_error_schema(["ERROR_TWO_FACTOR_AUTH_TYPE_DOES_NOT_EXIST"]),
+            429: get_error_schema(["ERROR_RATE_LIMIT_EXCEEDED"]),
         },
     )
     @map_exceptions(
         {
             TwoFactorAuthTypeDoesNotExist: ERROR_TWO_FACTOR_AUTH_TYPE_DOES_NOT_EXIST,
             VerificationFailed: ERROR_TWO_FACTOR_AUTH_VERIFICATION_FAILED,
+            RateLimitExceededException: ERROR_RATE_LIMIT_EXCEEDED,
         }
     )
     @validate_body(VerifyTOTPSerializer, return_validated=True)
@@ -204,7 +209,14 @@ def post(self, request, data: dict):
         Verifies TOTP two-factor authentication.
         """
 
-        TwoFactorAuthHandler().verify(TOTPAuthProviderType.type, **data)
+        def verify():
+            TwoFactorAuthHandler().verify(TOTPAuthProviderType.type, **data)
+
+        rate_limit(
+            rate=RateLimit.from_string("10/m"),
+            key=f"two_fa_verify:totp:{data.get('email', '')}",
+            raise_exception=True,
+        )(verify)()
 
         user = User.objects.filter(email=data["email"]).first()
         return_data = log_in_user(request, user)

backend/src/baserow/api/user/schemas.py

@@ -56,16 +56,37 @@
     }
 }
 
-create_user_response_schema = build_object_type(
+two_factor_required_response_schema = build_object_type(
+    {
+        "two_factor_auth": {
+            "type": "string",
+            "description": "The type of the two factor auth that is required to perform.",
+        },
+        "token": {
+            "type": "string",
+            "description": "The temporary token for verifying authentication using 2fa.",
+        },
+    }
+)
+
+success_create_user_response_schema = build_object_type(
     {
         **user_response_schema,
         **access_token_schema,
         **refresh_token_schema,
     }
 )
 
+authenticated_user_response_schema = {
+    "oneOf": [
+        success_create_user_response_schema,
+        two_factor_required_response_schema,
+    ],
+}
+
+
 if jwt_settings.ROTATE_REFRESH_TOKENS:
-    authenticate_user_schema = create_user_response_schema
+    authenticate_user_schema = authenticated_user_response_schema
 else:
     authenticate_user_schema = build_object_type(
         {

backend/src/baserow/api/user/serializers.py

@@ -307,6 +307,11 @@ def log_in_user(request, user):
     return data
 
 
+class TwoFactorAuthRequiredSerializer(serializers.Serializer):
+    two_factor_auth = serializers.CharField()
+    token = serializers.CharField()
+
+
 @extend_schema_serializer(deprecate_fields=["username"])
 class TokenObtainPairWithUserSerializer(TokenObtainPairSerializer):
     email = NormalizedEmailField(required=False)
@@ -343,10 +348,12 @@ def validate(self, attrs):
             if provider_type.is_enabled(twofa_provider):
                 token = TwoFactorAccessToken.for_user(self.user)
                 token.set_exp(lifetime=timedelta(minutes=2))
-                return {
-                    "two_factor_auth": provider_type.type,
-                    "token": str(token),
-                }
+                return TwoFactorAuthRequiredSerializer(
+                    {
+                        "two_factor_auth": provider_type.type,
+                        "token": str(token),
+                    }
+                ).data
 
         return log_in_user(self.context["request"], self.user)
 

backend/src/baserow/api/user/views.py

@@ -102,7 +102,7 @@
 from .exceptions import ClientSessionIdHeaderNotSetException
 from .schemas import (
     authenticate_user_schema,
-    create_user_response_schema,
+    authenticated_user_response_schema,
     verify_user_schema,
 )
 from .serializers import (
@@ -141,10 +141,12 @@ class ObtainJSONWebToken(TokenObtainPairView):
         operation_id="token_auth",
         description=(
             "Authenticates an existing user based on their email and their password. "
-            "If successful, an access token and a refresh token will be returned."
+            "If successful, an access token and a refresh token will be returned. "
+            "If the account is protected with two-factor authentication, "
+            "temporary token is returned to finish the verification."
         ),
         responses={
-            200: create_user_response_schema,
+            200: authenticated_user_response_schema,
             401: get_error_schema(
                 [
                     "ERROR_INVALID_CREDENTIALS",
@@ -269,7 +271,7 @@ class UserView(APIView):
             "account the initial workspace containing a database is created."
         ),
         responses={
-            200: create_user_response_schema,
+            200: authenticated_user_response_schema,
             400: get_error_schema(
                 [
                     "ERROR_ALREADY_EXISTS",
@@ -556,7 +558,7 @@ class VerifyEmailAddressView(APIView):
             "request is performed by unauthenticated user."
         ),
         responses={
-            200: create_user_response_schema,
+            200: authenticated_user_response_schema,
             400: get_error_schema(
                 [
                     "ERROR_INVALID_VERIFICATION_TOKEN",

backend/src/baserow/core/two_factor_auth/registries.py

@@ -3,6 +3,7 @@
 import string
 from abc import ABC, abstractmethod
 from base64 import b64encode
+from datetime import datetime, timedelta, timezone
 from io import BytesIO
 
 from django.conf import settings
@@ -117,6 +118,12 @@ def configure(
             raise TwoFactorAuthAlreadyConfigured
 
         if provider and kwargs.get("code"):
+            secret_valid_until = provider.created_on + timedelta(minutes=30)
+            now = datetime.now(tz=timezone.utc)
+            if now > secret_valid_until:
+                provider.delete()
+                raise VerificationFailed
+
             code = kwargs.get("code")
             totp = pyotp.TOTP(provider.secret)
 

backend/tests/baserow/api/users/test_token_auth.py

@@ -20,6 +20,7 @@
 from baserow.core.user.handler import UserHandler
 from baserow.core.user.utils import generate_session_tokens_for_user
 from baserow.core.utils import generate_hash
+from baserow.test_utils.helpers import AnyStr
 
 User = get_user_model()
 
@@ -278,6 +279,24 @@ def test_token_password_auth_disabled(api_client, data_fixture):
     }
 
 
+@pytest.mark.django_db
+def test_token_auth_2fa_required(api_client, data_fixture):
+    data_fixture.create_password_provider(enabled=True)
+    user, token = data_fixture.create_user_and_token(
+        email="test@localhost", password="test"
+    )
+    data_fixture.configure_totp(user)
+
+    response = api_client.post(
+        reverse("api:user:token_auth"),
+        {"email": "test@localhost", "password": "test"},
+        format="json",
+    )
+
+    assert response.status_code == HTTP_200_OK, response.json()
+    assert response.json() == {"two_factor_auth": "totp", "token": AnyStr()}
+
+
 @pytest.mark.django_db
 def test_token_password_auth_disabled_superadmin(api_client, data_fixture):
     data_fixture.create_password_provider(enabled=False)