From f205c48564aaceeff11ae3823cec97efb8bdff27 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felix=20Sp=C3=B6ttel?=
 <1682504+fspoettel@users.noreply.github.com>
Date: Wed, 2 Sep 2020 19:12:15 +0200
Subject: [PATCH] Use pyotp for validating TOTP codes

* also implements resynchronisation support via `pyotp`'s `valid_window option
---
 management/totp.py  | 30 ++++++------------------------
 setup/management.sh |  2 +-
 2 files changed, 7 insertions(+), 25 deletions(-)

diff --git a/management/totp.py b/management/totp.py
index 52cdc256..908b4764 100644
--- a/management/totp.py
+++ b/management/totp.py
@@ -4,20 +4,16 @@ import io
 import os
 import struct
 import time
-from urllib.parse import quote
+import pyotp
 import qrcode
 
 def get_secret():
 	return base64.b32encode(os.urandom(20)).decode('utf-8')
 
 def get_otp_uri(secret, email):
-	site_name = 'mailinabox'
-
-	return 'otpauth://totp/{}:{}?secret={}&issuer={}'.format(
-		quote(site_name),
-		quote(email),
-		secret,
-		quote(site_name)
+	return pyotp.TOTP(secret).provisioning_uri(
+		name=email,
+		issuer_name='mailinabox'
 	)
 
 def get_qr_code(data):
@@ -32,20 +28,6 @@ def validate(secret, token):
 	"""
 	@see https://tools.ietf.org/html/rfc6238#section-4
 	@see https://tools.ietf.org/html/rfc4226#section-5.4
-	@see https://git.sr.ht/~sircmpwn/meta.sr.ht/tree/master/metasrht/totp.py
-	@see https://github.com/susam/mintotp/blob/master/mintotp.py
-	TODO: resynchronisation
 	"""
-	key = base64.b32decode(secret)
-	tm = int(time.time() / 30)
-	digits = 6
-
-	step = 0
-	counter = struct.pack('>Q', tm + step)
-
-	hm = hmac.HMAC(key, counter, 'sha1').digest()
-	offset = hm[-1] &0x0F
-	binary = struct.unpack(">L", hm[offset:offset + 4])[0] & 0x7fffffff
-
-	code = str(binary)[-digits:].rjust(digits, '0')
-	return token == code
+	totp = pyotp.TOTP(secret)
+	return totp.verify(token, valid_window=2)
diff --git a/setup/management.sh b/setup/management.sh
index ce78b171..ae942985 100755
--- a/setup/management.sh
+++ b/setup/management.sh
@@ -50,7 +50,7 @@ hide_output $venv/bin/pip install --upgrade pip
 hide_output $venv/bin/pip install --upgrade \
 	rtyaml "email_validator>=1.0.0" "exclusiveprocess" \
 	flask dnspython python-dateutil \
-    qrcode[pil] \
+    qrcode[pil] pyotp \
 	"idna>=2.0.0" "cryptography==2.2.2" boto psutil postfix-mta-sts-resolver
 
 # CONFIGURATION