From 694e7ddd7cb0d942132387b6c4cb0850fccb14b0 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Sun, 16 Jan 2022 16:54:13 +0800
Subject: [PATCH] =?UTF-8?q?CAPTCHA=E5=A2=9E=E5=8A=A0=E5=A4=9A=E4=B8=AA?=
 =?UTF-8?q?=E9=80=89=E9=A1=B9?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/waf/action_captcha.go    | 19 ++++++++++++-------
 internal/waf/captcha_validator.go | 18 ++++++++++++++++--
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/internal/waf/action_captcha.go b/internal/waf/action_captcha.go
index 8356ea1..14147f2 100644
--- a/internal/waf/action_captcha.go
+++ b/internal/waf/action_captcha.go
@@ -18,7 +18,10 @@ const (
 )
 
 type CaptchaAction struct {
-	Life           int32  `yaml:"life" json:"life"`
+	Life             int32 `yaml:"life" json:"life"`
+	MaxFails         int   `yaml:"maxFails" json:"maxFails"`                 // 最大失败次数
+	FailBlockTimeout int   `yaml:"failBlockTimeout" json:"failBlockTimeout"` // 失败拦截时间
+
 	Language       string `yaml:"language" json:"language"`             // 语言，zh-CN, en-US ...
 	AddToWhiteList bool   `yaml:"addToWhiteList" json:"addToWhiteList"` // 是否加入到白名单
 	Scope          string `yaml:"scope" json:"scope"`
@@ -60,12 +63,14 @@ func (this *CaptchaAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, req
 	}
 
 	var captchaConfig = maps.Map{
-		"action":    this,
-		"timestamp": time.Now().Unix(),
-		"url":       refURL,
-		"policyId":  waf.Id,
-		"groupId":   group.Id,
-		"setId":     set.Id,
+		"action":           this,
+		"timestamp":        time.Now().Unix(),
+		"maxFails":         this.MaxFails,
+		"failBlockTimeout": this.FailBlockTimeout,
+		"url":              refURL,
+		"policyId":         waf.Id,
+		"groupId":          group.Id,
+		"setId":            set.Id,
 	}
 	info, err := utils.SimpleEncryptMap(captchaConfig)
 	if err != nil {
diff --git a/internal/waf/captcha_validator.go b/internal/waf/captcha_validator.go
index 775db5b..3d1fee2 100644
--- a/internal/waf/captcha_validator.go
+++ b/internal/waf/captcha_validator.go
@@ -3,6 +3,8 @@ package waf
 import (
 	"bytes"
 	"encoding/base64"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
+	"github.com/TeaOSLab/EdgeNode/internal/ttlcache"
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
 	"github.com/TeaOSLab/EdgeNode/internal/utils/jsonutils"
 	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
@@ -54,7 +56,7 @@ func (this *CaptchaValidator) Run(request requests.Request, writer http.Response
 	var originURL = m.GetString("url")
 
 	if request.WAFRaw().Method == http.MethodPost && len(request.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_ID")) > 0 {
-		this.validate(actionConfig, m.GetInt64("policyId"), m.GetInt64("groupId"), setId, originURL, request, writer)
+		this.validate(actionConfig, m.GetInt("maxFails"), m.GetInt("failBlockTimeout"), m.GetInt64("policyId"), m.GetInt64("groupId"), setId, originURL, request, writer)
 	} else {
 		this.show(actionConfig, request, writer)
 	}
@@ -142,11 +144,14 @@ func (this *CaptchaValidator) show(actionConfig *CaptchaAction, request requests
 </html>`))
 }
 
-func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, policyId int64, groupId int64, setId int64, originURL string, request requests.Request, writer http.ResponseWriter) (allow bool) {
+func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, maxFails int, failBlockTimeout int, policyId int64, groupId int64, setId int64, originURL string, request requests.Request, writer http.ResponseWriter) (allow bool) {
 	captchaId := request.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_ID")
 	if len(captchaId) > 0 {
 		captchaCode := request.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_CODE")
 		if captcha.VerifyString(captchaId, captchaCode) {
+			// 删除计数
+			ttlcache.SharedCache.Delete("CAPTCHA:FAILS:" + request.WAFRemoteIP())
+
 			var life = CaptchaSeconds
 			if actionConfig.Life > 0 {
 				life = types.Int(actionConfig.Life)
@@ -159,6 +164,15 @@ func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, policyId int
 
 			return false
 		} else {
+			// 增加计数
+			if maxFails > 0 && failBlockTimeout > 0 {
+				var countFails = ttlcache.SharedCache.IncreaseInt64("CAPTCHA:FAILS:"+request.WAFRemoteIP(), 1, time.Now().Unix()+300)
+				if int(countFails) >= maxFails {
+					SharedIPBlackList.RecordIP(IPTypeAll, firewallconfigs.FirewallScopeService, request.WAFServerId(), request.WAFRemoteIP(), time.Now().Unix()+int64(failBlockTimeout), policyId, false, groupId, setId, "CAPTCHA验证连续失败")
+					return false
+				}
+			}
+
 			http.Redirect(writer, request.WAFRaw(), request.WAFRaw().URL.String(), http.StatusSeeOther)
 		}
 	}