Fixed #60 -- Always cast the token to an int before verification

jazzband · Jun 4, 2014 · 02ee37b · 02ee37b
1 parent 4b3c664
commit 02ee37b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/tests/tests.py b/tests/tests.py
@@ -792,6 +792,15 @@ def test_verify(self):
         self.assertFalse(device.verify_token(-1))
         self.assertTrue(device.verify_token(totp(device.bin_key)))
 
+    def test_verify_token_as_string(self):
+        """
+        The field used to read the token may be a CharField,
+        so the PhoneDevice must be able to validate tokens
+        read as strings
+        """
+        device = PhoneDevice(key=random_hex().decode())
+        self.assertTrue(device.verify_token(str(totp(device.bin_key))))
+
     def test_unicode(self):
         device = PhoneDevice(name='unknown')
         self.assertEqual('unknown (None)', str(device))

diff --git a/two_factor/models.py b/two_factor/models.py
@@ -89,6 +89,11 @@ def bin_key(self):
         return unhexlify(self.key.encode())
 
     def verify_token(self, token):
+        try:
+            token = int(token)
+        except ValueError:
+            return False
+
         for drift in range(-5, 1):
             if totp(self.bin_key, drift=drift) == token:
                 return True