Add support for TOTP in Python wrapper

2 files changed, 39 insertions, 1 deletions

diff --git a/src/python/botan2.py b/src/python/botan2.py
index e7e685802..7db8cd9b3 100755
--- a/src/python/botan2.py
+++ b/src/python/botan2.py
@@ -22,7 +22,7 @@ from ctypes import CDLL, POINTER, byref, create_string_buffer, \
     c_void_p, c_size_t, c_uint8, c_uint32, c_uint64, c_int, c_uint, c_char_p
 
 from sys import version_info, platform
-from time import strptime, mktime
+from time import strptime, mktime, time as system_time
 from binascii import hexlify
 from datetime import datetime
 
@@ -1588,6 +1588,29 @@ class HOTP(object):
         else:
             return (False, counter)
 
+class TOTP(object):
+    def __init__(self, key, digest="SHA-1", digits=6, timestep=30):
+        self.__obj = c_void_p(0)
+        _DLL.botan_totp_init(byref(self.__obj), key, len(key), _ctype_str(digest), digits, timestep)
+
+    def __del__(self):
+        _DLL.botan_totp_destroy(self.__obj)
+
+    def generate(self, timestamp=None):
+        if timestamp is None:
+            timestamp = int(system_time())
+        code = c_uint32(0)
+        _DLL.botan_totp_generate(self.__obj, byref(code), timestamp)
+        return code.value
+
+    def check(self, code, timestamp=None, acceptable_drift=0):
+        if timestamp is None:
+            timestamp = int(system_time())
+        rc = _DLL.botan_totp_check(self.__obj, code, timestamp, acceptable_drift)
+        if rc == 0:
+            return True
+        return False
+
 def nist_key_wrap(kek, key):
     output = create_string_buffer(len(key) + 8)
     out_len = c_size_t(len(output))
diff --git a/src/scripts/test_python.py b/src/scripts/test_python.py
index 04b0c949e..e860df358 100644
--- a/src/scripts/test_python.py
+++ b/src/scripts/test_python.py
@@ -596,6 +596,21 @@ ofvkP1EDmpx50fHLawIDAQAB
         self.assertEqual(hotp.check(520489, 7, 2), (True, 10))
         self.assertEqual(hotp.check(520489, 0, 9), (True, 10))
 
+    def test_totp(self):
+
+        totp = botan2.TOTP(b'12345678901234567890', digest="SHA-1", digits=8)
+
+        self.assertEqual(totp.generate(59), 94287082)
+        self.assertEqual(totp.generate(1111111109), 7081804)
+        self.assertEqual(totp.generate(1111111111), 14050471)
+        self.assertEqual(totp.generate(1234567890), 89005924)
+        self.assertEqual(totp.generate(1234567890), 89005924)
+        self.assertEqual(totp.generate(2000000000), 69279037)
+
+        self.assertTrue(totp.check(7081804, 1111111109))
+        self.assertTrue(totp.check(7081804, 1111111109 - 29))
+        self.assertFalse(totp.check(7081804, 1111111109 + 1))
+        self.assertTrue(totp.check(7081804, 1111111109 + 30, 1))
 
 if __name__ == '__main__':
     unittest.main()