From 5299d87c82ace2e468eef9811d79fd9665064118 Mon Sep 17 00:00:00 2001
From: Daniel Townsend <dan@dantownsend.co.uk>
Date: Wed, 21 Aug 2024 00:14:07 +0100
Subject: [PATCH] fix rate limit for MFA setup endpoint

---
 piccolo_admin/endpoints.py | 32 +++++++++++++++++++-------------
 piccolo_admin/example.py   | 10 ++++++----
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/piccolo_admin/endpoints.py b/piccolo_admin/endpoints.py
index 1811057c..648d6012 100644
--- 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
index afada616..922cf1cf 100644
--- 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,
+        ),
+    ],
 )