WAF增加多个动作

internal/waf/action_captcha.go

@@ -1,11 +1,14 @@
 package waf
 
 import (
+	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
+	"github.com/TeaOSLab/EdgeNode/internal/utils"
 	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
-	"github.com/iwind/TeaGo/types"
+	"github.com/iwind/TeaGo/maps"
 	stringutil "github.com/iwind/TeaGo/utils/string"
 	"net/http"
 	"net/url"
+	"strings"
 	"time"
 )
 
@@ -13,27 +16,63 @@ var captchaSalt = stringutil.Rand(32)
 
 const (
 	CaptchaSeconds = 600 // 10 minutes
+	CaptchaPath    = "/WAF/VERIFY/CAPTCHA"
 )
 
 type CaptchaAction struct {
+	Life           int32  `yaml:"life" json:"life"`
+	Language       string `yaml:"language" json:"language"`             // 语言，zh-CN, en-US ...
+	AddToWhiteList bool   `yaml:"addToWhiteList" json:"addToWhiteList"` // 是否加入到白名单
 }
 
-func (this *CaptchaAction) Perform(waf *WAF, request *requests.Request, writer http.ResponseWriter) (allow bool) {
-	// TEAWEB_CAPTCHA:
-	cookie, err := request.Cookie("TEAWEB_WAF_CAPTCHA")
-	if err == nil && cookie != nil && len(cookie.Value) > 32 {
-		m := cookie.Value[:32]
-		timestamp := cookie.Value[32:]
-		if stringutil.Md5(captchaSalt+timestamp) == m && time.Now().Unix() < types.Int64(timestamp) { // verify md5
-			return true
+func (this *CaptchaAction) Init(waf *WAF) error {
+	return nil
+}
+
+func (this *CaptchaAction) Code() string {
+	return ActionCaptcha
+}
+
+func (this *CaptchaAction) IsAttack() bool {
+	return false
+}
+
+func (this *CaptchaAction) WillChange() bool {
+	return true
+}
+
+func (this *CaptchaAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) (allow bool) {
+	// 是否在白名单中
+	if SharedIPWhiteList.Contains("set:"+set.Id, request.WAFRemoteIP()) {
+		return true
+	}
+
+	refURL := request.WAFRaw().URL.String()
+
+	// 覆盖配置
+	if strings.HasPrefix(refURL, CaptchaPath) {
+		info := request.WAFRaw().URL.Query().Get("info")
+		if len(info) > 0 {
+			m, err := utils.SimpleDecryptMap(info)
+			if err == nil && m != nil {
+				refURL = m.GetString("url")
+			}
 		}
 	}
 
-	refURL := request.URL.String()
-	if len(request.Referer()) > 0 {
-		refURL = request.Referer()
+	var captchaConfig = maps.Map{
+		"action":    this,
+		"timestamp": time.Now().Unix(),
+		"url":       refURL,
+		"setId":     set.Id,
 	}
-	http.Redirect(writer, request.Raw(), "/WAFCAPTCHA?url="+url.QueryEscape(refURL), http.StatusTemporaryRedirect)
+	info, err := utils.SimpleEncryptMap(captchaConfig)
+	if err != nil {
+		remotelogs.Error("WAF_CAPTCHA_ACTION", "encode captcha config failed: "+err.Error())
+		return true
+	}
+
+	http.Redirect(writer, request.WAFRaw(), CaptchaPath+"?info="+url.QueryEscape(info), http.StatusTemporaryRedirect)
 
 	return false
 }