use calcOTP() wrapper

2 files changed, 4 insertions, 37 deletions

diff --git a/game/update.rpy b/game/update.rpy
index 7b9f555..6f24134 100644
--- a/game/update.rpy
+++ b/game/update.rpy
@@ -383,22 +383,11 @@ label register:
     if persistent.autologin and persistent.email and persistent.passd and persistent.totp:
         python:
             print("Automatic login ON")
-            key = base64.b32decode(persistent.totp.encode('utf-8'), True)
-            msg = struct.pack(">Q", int(time.time()/30))
-            h = hmac.new(key, msg, hashlib.sha1).digest()
-            if LEGACY:
-                o = ord(h[19]) & 15
-            else:
-                o = (h[19] & 15)
-            _return = (struct.unpack(">I", h[o:o+4])[0] & 0x7fffffff) % 1000000
-            _return = "%06d" % _return
-            print("TOTP: %s" % _return)
-            del key, msg, h, o
+            code2FA = calcOTP(base64.b32decode(persistent.totp.encode('utf-8'), True))
             if LEGACY:
                 password = str(bytearray((x ^ int(persistent.rhash/mp.sub) for x in bytearray(persistent.passd, 'utf-8'))))
             else:
                 password = bytearray((x ^ int(persistent.rhash/mp.sub) for x in persistent.passd)).decode('utf-8')
-            code2FA = _return
             data = {"mail": persistent.email,
                     "pass": password,
                     "totp": code2FA[:6]
@@ -485,20 +474,9 @@ label register:
         $ status_update(pc=92)
         if persistent.totp is None:
             call screen register_input(_("If you already have an account, please insert your {b}2FA code{/b}.\n\n{u}Otherwise, a new account will be created and details will be sent to your email.{/u}"))
+            $ code2FA = _return
         else:
-            python:
-                key = base64.b32decode(persistent.totp.encode('utf-8'), True)
-                msg = struct.pack(">Q", int(time.time()/30))
-                h = hmac.new(key, msg, hashlib.sha1).digest()
-                if LEGACY:
-                    o = ord(h[19]) & 15
-                else:
-                    o = (h[19] & 15)
-                _return = (struct.unpack(">I", h[o:o+4])[0] & 0x7fffffff) % 1000000
-                _return = "%06d" % _return
-                print("TOTP: %s" % _return)
-                del key, msg, h, o
-        $ code2FA = _return
+            $ code2FA = calcOTP(base64.b32decode(persistent.totp.encode('utf-8'), True))
 
         $ status_update(pc=95)
         $ data = {"mail": email,
diff --git a/game/vault.rpy b/game/vault.rpy
index e421e14..ef9634e 100644
--- a/game/vault.rpy
+++ b/game/vault.rpy
@@ -181,18 +181,7 @@ label register_vault:
     $ status_update(pc=92)
     if persistent.totp is not None:
         python:
-            key = base64.b32decode(persistent.totp.encode('utf-8'), True)
-            msg = struct.pack(">Q", int(time.time()/30))
-            h = hmac.new(key, msg, hashlib.sha1).digest()
-            if LEGACY:
-                o = ord(h[19]) & 15
-            else:
-                o = (h[19] & 15)
-            _return = (struct.unpack(">I", h[o:o+4])[0] & 0x7fffffff) % 1000000
-            _return = "%06d" % _return
-            print("TOTP: %s" % _return)
-            uedit["totp"] = _return
-            del key, msg, h, o
+            uedit["totp"] = calcOTP(base64.b32decode(persistent.totp.encode('utf-8'), True))
 
     $ email    = uedit["mail"]
     $ password = uedit["pasd"]