fix rate limit for MFA setup endpoint

piccolo-orm · Aug 20, 2024 · 5299d87 · 5299d87
1 parent 6672ef9
commit 5299d87
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 17 deletions.
diff --git a/piccolo_admin/endpoints.py b/piccolo_admin/endpoints.py
@@ -433,7 +433,7 @@ def __init__(
         allowed_hosts: t.Sequence[str] = [],
         debug: bool = False,
         sidebar_links: t.Dict[str, str] = {},
-        mfa_provider: t.Optional[MFAProvider] = None,
+        mfa_providers: t.Optional[t.Sequence[MFAProvider]] = None,
     ) -> None:
         super().__init__(
             title=site_name,
@@ -690,17 +690,23 @@ def __init__(
         #######################################################################
         # MFA
 
-        if mfa_provider:
-            private_app.mount(
-                path="/mfa-setup/",
-                app=RateLimitingMiddleware(
-                    app=mfa_setup(
-                        provider=mfa_provider,
-                        auth_table=self.auth_table,
+        if mfa_providers:
+            if len(mfa_providers) > 1:
+                raise ValueError(
+                    "Only a single mfa_provider is currently supported."
+                )
+
+            for mfa_provider in mfa_providers:
+                private_app.mount(
+                    path="/mfa-setup/",
+                    app=RateLimitingMiddleware(
+                        app=mfa_setup(
+                            provider=mfa_provider,
+                            auth_table=self.auth_table,
+                        ),
+                        provider=InMemoryLimitProvider(limit=5, timespan=300),
                     ),
-                    provider=rate_limit_provider,
-                ),
-            )
+                )
 
         #######################################################################
 
@@ -1106,7 +1112,7 @@ def create_admin(
     allowed_hosts: t.Sequence[str] = [],
     debug: bool = False,
     sidebar_links: t.Dict[str, str] = {},
-    mfa_provider: t.Optional[MFAProvider] = None,
+    mfa_providers: t.Optional[t.Sequence[MFAProvider]] = None,
 ):
     """
     :param tables:
@@ -1273,5 +1279,5 @@ def create_admin(
         allowed_hosts=allowed_hosts,
         debug=debug,
         sidebar_links=sidebar_links,
-        mfa_provider=mfa_provider,
+        mfa_providers=mfa_providers,
     )
diff --git a/piccolo_admin/example.py b/piccolo_admin/example.py
@@ -616,10 +616,12 @@ def booking_endpoint(request: Request, data: BookingModel) -> str:
         "Top Movies": "/#/movie?__order=-box_office",
         "Google": "https://google.com",
     },
-    mfa_provider=AuthenticatorProvider(
-        db_encryption_key="wqsOqyTTEsrWppZeIMS8a3l90yPUtrqT48z7FS6_U8g=",
-        secret_table=AuthenticatorSecret,
-    ),
+    mfa_providers=[
+        AuthenticatorProvider(
+            db_encryption_key="wqsOqyTTEsrWppZeIMS8a3l90yPUtrqT48z7FS6_U8g=",
+            secret_table=AuthenticatorSecret,
+        ),
+    ],
 )