From ab6c88823d39f240fcdb3d9c01f3284c3f75e7b9 Mon Sep 17 00:00:00 2001
From: vkill <vkill.net@gmail.com>
Date: Mon, 17 Dec 2018 14:26:00 +0800
Subject: [PATCH] Support for TOTP valid_window configuration (#2187)

---
 apps/jumpserver/settings.py | 1 +
 apps/users/utils.py         | 3 ++-
 config_docker.py            | 7 +++++++
 config_example.py           | 3 +++
 4 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/apps/jumpserver/settings.py b/apps/jumpserver/settings.py
index fb595a4df..a58642877 100644
--- a/apps/jumpserver/settings.py
+++ b/apps/jumpserver/settings.py
@@ -356,6 +356,7 @@ FILE_UPLOAD_DIRECTORY_PERMISSIONS = 0o755
 
 # OTP settings
 OTP_ISSUER_NAME = CONFIG.OTP_ISSUER_NAME
+OTP_VALID_WINDOW = CONFIG.OTP_VALID_WINDOW
 
 # Auth LDAP settings
 AUTH_LDAP = False
diff --git a/apps/users/utils.py b/apps/users/utils.py
index eac1c6f99..c998774b0 100644
--- a/apps/users/utils.py
+++ b/apps/users/utils.py
@@ -292,7 +292,8 @@ def check_otp_code(otp_secret_key, otp_code):
     if not otp_secret_key or not otp_code:
         return False
     totp = pyotp.TOTP(otp_secret_key)
-    return totp.verify(otp_code)
+    otp_valid_window = settings.OTP_VALID_WINDOW or 0
+    return totp.verify(otp=otp_code, valid_window=otp_valid_window)
 
 
 def get_password_check_rules():
diff --git a/config_docker.py b/config_docker.py
index ca322b4b8..643c11055 100644
--- a/config_docker.py
+++ b/config_docker.py
@@ -100,6 +100,9 @@ class Config:
     }
     AUTH_LDAP_START_TLS = False
 
+    #
+    # OTP_VALID_WINDOW = 0
+
     def __init__(self):
         pass
 
@@ -200,6 +203,10 @@ class DockerConfig(Config):
     AUTH_LDAP_START_TLS = False
 
 
+    #
+    OTP_VALID_WINDOW = int(os.environ.get("OTP_VALID_WINDOW")) if os.environ.get("OTP_VALID_WINDOW") else 0
+
+
 # Default using Config settings, you can write if/else for different env
 config = DockerConfig()
 
diff --git a/config_example.py b/config_example.py
index dfcc876a3..e37df23b0 100644
--- a/config_example.py
+++ b/config_example.py
@@ -90,6 +90,9 @@ class Config:
     # AUTH_OPENID_CLIENT_ID = 'client-id'
     # AUTH_OPENID_CLIENT_SECRET = 'client-secret'
 
+    #
+    # OTP_VALID_WINDOW = 0
+
     def __init__(self):
         pass