From 3d60ab1fc53245644b7aa933fae48ff11baa0829 Mon Sep 17 00:00:00 2001
From: samwaf <samwafgo@gmail.com>
Date: Wed, 5 Mar 2025 09:58:16 +0800
Subject: [PATCH] feat: otp force

#169
---
 global/config.go                  |  1 +
 middleware/auth_api_check.go      | 11 +++++++++++
 model/common/response/response.go |  4 ++++
 waftask/task_config.go            |  4 ++++
 4 files changed, 20 insertions(+)

diff --git a/global/config.go b/global/config.go
index d4e49e3..d0890d2 100644
--- a/global/config.go
+++ b/global/config.go
@@ -30,6 +30,7 @@ var (
 	GCONFIG_RECORD_DNS_NORMAL_EXPIRE_HOURS int64  = 7 * 24 //DNS 正常有效期 单位小时 默认7天
 	GCONFIG_RECORD_SPIDER_DENY             int64  = 0      //爬虫禁止访问开关 默认 0 只检测不阻止访问 1 检测并阻止访问
 	GCONFIG_RECORD_HIDE_SERVER_HEADER      int64  = 1      // 是否隐藏Server头信息 1隐藏 0不隐藏
+	GCONFIG_RECORD_FORCE_BIND_2FA          int64  = 0      // 是否强制绑定双因素认证(1强制 0不强制)
 	GCONFIG_RECORD_DEBUG_ENABLE            int64  = 0      //调试开关 默认关闭
 	GCONFIG_RECORD_DEBUG_PWD               string = ""     //调试密码 如果未空则不需要密码
 
diff --git a/middleware/auth_api_check.go b/middleware/auth_api_check.go
index 75e090e..4a9a4d4 100644
--- a/middleware/auth_api_check.go
+++ b/middleware/auth_api_check.go
@@ -15,6 +15,7 @@ import (
 
 var (
 	wafTokenInfoService = waf_service.WafTokenInfoServiceApp
+	wafOtpService       = waf_service.WafOtpServiceApp
 )
 
 // Auth 鉴权中间件
@@ -70,6 +71,16 @@ func Auth() gin.HandlerFunc {
 						}
 					}
 
+					//检测是否强制2Fa绑定
+					if global.GCONFIG_RECORD_FORCE_BIND_2FA == 1 && c.Request.RequestURI != "/samwaf/ws" && c.Request.RequestURI != "/samwaf/logout" {
+						otpBean := wafOtpService.GetDetailByUserNameApi(tokenInfo.LoginAccount)
+						if otpBean.UserName == "" {
+							//需要强制跳转2fa绑定界面
+							response.NeedBind2FAWithMessage("系统已开启强制 【双因素认证】 ，请进行绑定", c)
+							c.Abort()
+							return
+						}
+					}
 				}
 			}
 		}
diff --git a/model/common/response/response.go b/model/common/response/response.go
index 27edc3f..337ec58 100644
--- a/model/common/response/response.go
+++ b/model/common/response/response.go
@@ -18,6 +18,7 @@ const (
 	ERROR             = -1
 	SUCCESS           = 0
 	INPUT_SECRET_CODE = -2
+	NEED_BIND_2FA     = -3
 	AUTHFAIL          = -999
 )
 
@@ -65,3 +66,6 @@ func AuthFailWithMessage(message string, c *gin.Context) {
 func SecretCodeFailWithMessage(message string, c *gin.Context) {
 	Result(INPUT_SECRET_CODE, map[string]interface{}{}, message, c)
 }
+func NeedBind2FAWithMessage(message string, c *gin.Context) {
+	Result(NEED_BIND_2FA, map[string]interface{}{}, message, c)
+}
diff --git a/waftask/task_config.go b/waftask/task_config.go
index ac301f6..152ba8f 100644
--- a/waftask/task_config.go
+++ b/waftask/task_config.go
@@ -76,6 +76,9 @@ func setConfigIntValue(name string, value int64, change int) {
 	case "hide_server_header":
 		global.GCONFIG_RECORD_HIDE_SERVER_HEADER = value
 		break
+	case "force_bind_2fa":
+		global.GCONFIG_RECORD_FORCE_BIND_2FA = value
+		break
 	default:
 		zlog.Warn("Unknown config item:", name)
 	}
@@ -214,4 +217,5 @@ func TaskLoadSetting(initLoad bool) {
 	updateConfigStringItem(initLoad, "gpt", "gpt_token", global.GCONFIG_RECORD_GPT_TOKEN, "GPT远程授权密钥", "string", "")
 	updateConfigStringItem(initLoad, "gpt", "gpt_model", global.GCONFIG_RECORD_GPT_MODEL, "GPT模型名称", "string", "")
 	updateConfigIntItem(initLoad, "security", "hide_server_header", global.GCONFIG_RECORD_HIDE_SERVER_HEADER, "是否隐藏Server响应头(1隐藏 0不隐藏)", "int", "")
+	updateConfigIntItem(initLoad, "security", "force_bind_2fa", global.GCONFIG_RECORD_FORCE_BIND_2FA, "是否强制绑定双因素认证(1强制 0不强制)", "options", "0|不强制,1|强制")
 }