implement a method for TOTP that returns the matching timecode

Author: jap

src/pyotp/totp.py

@@ -2,7 +2,7 @@
 import datetime
 import hashlib
 import time
-from typing import Any, Optional, Union
+from typing import Any, Literal, Optional, Union
 
 from . import utils
 from .otp import OTP
@@ -40,7 +40,7 @@ def __init__(
 
     def at(self, for_time: Union[int, datetime.datetime], counter_offset: int = 0) -> str:
         """
-        Accepts either a Unix timestamp integer or a datetime object.
+        Accepts either a Unix timestamp or a datetime object.
 
         To get the time until the next timecode change (seconds until the current OTP expires), use this instead:
 
@@ -85,6 +85,35 @@ def verify(self, otp: str, for_time: Optional[datetime.datetime] = None, valid_w
 
         return utils.strings_equal(str(otp), str(self.at(for_time)))
 
+    def verify_and_get_timecode(
+        self, otp: str, for_time: Optional[Union[int, datetime.datetime]] = None, valid_window: int = 0
+    ) -> int | Literal[False]:
+        """
+        Verifies the OTP passed in against the current time OTP and returns the matching timecode.
+
+        This is useful when callers need to prevent TOTP reuse. An
+        application can store the timecode from the last successful
+        authentication and reject future OTPs whose matching timecode
+        is not greater than the stored one.
+
+        :param otp: the OTP to check against
+        :param for_time: Time to check OTP at (defaults to now); accepts datetime or Unix timestamp
+        :param valid_window: extends the validity to this counter ticks before and after the current one
+        :returns: matching timecode if verification succeeded, False otherwise
+        """
+        if for_time is None:
+            for_time = datetime.datetime.now()
+        elif isinstance(for_time, (float, int)):
+            for_time = datetime.datetime.fromtimestamp(int(for_time))
+
+        if valid_window < 0:
+            raise ValueError("valid_window cannot be negative")
+
+        for i in range(-valid_window, valid_window + 1):
+            if utils.strings_equal(str(otp), str(self.at(for_time, i))):
+                return self.timecode(for_time) + i
+        return False
+
     def provisioning_uri(self, name: Optional[str] = None, issuer_name: Optional[str] = None, **kwargs) -> str:
         """
         Returns the provisioning URI for the OTP.  This can then be

test.py

@@ -335,6 +335,15 @@ def test_valid_window(self):
         self.assertFalse(totp.verify("195979", 200, 1))
 
 
+class GetTimeCodeTest(unittest.TestCase):
+    def test_timecode(self):
+        totp = pyotp.TOTP("ABCDEFGH")
+        self.assertEqual(totp.verify_and_get_timecode("451564", 200, 1), 5)
+        self.assertEqual(totp.verify_and_get_timecode("028307", 200, 1), 6)
+        self.assertEqual(totp.verify_and_get_timecode("681610", 200, 1), 7)
+        self.assertFalse(totp.verify_and_get_timecode("195979", 200, 1))
+
+
 class DigestFunctionTest(unittest.TestCase):
     def test_md5(self):
         with self.assertRaises(ValueError) as cm: