From 01b7b3684860f45e7d17665a75033c2ec548540c Mon Sep 17 00:00:00 2001
From: Daniel Townsend <dan@dantownsend.co.uk>
Date: Sun, 8 Sep 2024 13:44:50 +0100
Subject: [PATCH] add `valid_window` argument to `AuthenticatorProvider`

---
 piccolo_api/mfa/authenticator/provider.py |  7 +++++++
 piccolo_api/mfa/authenticator/tables.py   | 14 ++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/piccolo_api/mfa/authenticator/provider.py b/piccolo_api/mfa/authenticator/provider.py
index f8fca9a..dadcfa8 100644
--- a/piccolo_api/mfa/authenticator/provider.py
+++ b/piccolo_api/mfa/authenticator/provider.py
@@ -27,6 +27,7 @@ def __init__(
         issuer_name: str = "Piccolo-MFA",
         register_template_path: t.Optional[str] = None,
         styles: t.Optional[Styles] = None,
+        valid_window: int = 0,
     ):
         """
         Allows authentication using an authenticator app on the user's phone,
@@ -51,6 +52,10 @@ def __init__(
             visual changes.
         :param styles:
             Modify the appearance of the HTML template using CSS.
+        :param valid_window:
+            Extends the validity to this many counter ticks before and after
+            the current one. Increasing it is more convenient for users, but
+            is less secure.
 
         """
         super().__init__(
@@ -62,6 +67,7 @@ def __init__(
         self.secret_table = secret_table
         self.issuer_name = issuer_name
         self.styles = styles or Styles()
+        self.valid_window = valid_window
 
         # Load the Jinja Template
         register_template_path = (
@@ -81,6 +87,7 @@ async def authenticate_user(self, user: BaseUser, code: str) -> bool:
             user_id=user.id,
             code=code,
             encryption_provider=self.encryption_provider,
+            valid_window=self.valid_window,
         )
 
     async def is_user_enrolled(self, user: BaseUser) -> bool:
diff --git a/piccolo_api/mfa/authenticator/tables.py b/piccolo_api/mfa/authenticator/tables.py
index ebfe93a..466ac2a 100644
--- a/piccolo_api/mfa/authenticator/tables.py
+++ b/piccolo_api/mfa/authenticator/tables.py
@@ -138,8 +138,18 @@ async def revoke(cls, user_id: int):
 
     @classmethod
     async def authenticate(
-        cls, user_id: int, code: str, encryption_provider: EncryptionProvider
+        cls,
+        user_id: int,
+        code: str,
+        encryption_provider: EncryptionProvider,
+        valid_window: int = 0,
     ) -> bool:
+        """
+        :param valid_window:
+            Extends the validity to this many counter ticks before and after
+            the current one.
+
+        """
         secret = (
             await cls.objects()
             .where(
@@ -166,7 +176,7 @@ async def authenticate(
         )
         totp = pyotp.TOTP(shared_secret)  # type: ignore
 
-        if totp.verify(code):
+        if totp.verify(code, valid_window=valid_window):
             secret.last_used_at = datetime.datetime.now(
                 tz=datetime.timezone.utc
             )