feat: check ip response error count

#529

enums/cache_enum.go

@@ -11,5 +11,6 @@ const (
 	CACHE_CAPTCHA_PASS   = "CACHE_CAPTCHA_PASS"    //通过验证的码
 	CACHE_ANNOUNCEMENT   = "CACHE_ANNOUNCEMENT"    //公告数据
 	CACHE_WEBFILE        = "CACHE_WEBFILE"
-	CACHE_FILE_INFO      = "CACHE_FILE_INFO" //文件信息
+	CACHE_FILE_INFO      = "CACHE_FILE_INFO"      //文件信息
+	CACHE_IP_FAILURE_PRE = "CACHE_IP_FAILURE_PRE" //IP失败记录前缀
 )

global/config.go

@@ -54,4 +54,10 @@ var (
 	GCONFIG_RECORD_GPT_TOKEN string = "SamWaf提示请输入密钥"            //GPT远程授权密钥
 	GCONFIG_RECORD_GPT_MODEL string = "deepseek-chat"            //GPT 模型名称
 
+	// IP失败封禁相关配置
+	GCONFIG_IP_FAILURE_STATUS_CODES    string = "401|403|404|444|429|503" //失败状态码配置，支持多个用|分隔，也支持正则表达式
+	GCONFIG_IP_FAILURE_BAN_ENABLED     int64  = 0                         //是否启用IP失败封禁 1启用 0禁用
+	GCONFIG_IP_FAILURE_BAN_TIME_WINDOW int64  = 5                         //IP失败封禁时间窗口（分钟）默认5分钟
+	GCONFIG_IP_FAILURE_BAN_MAX_COUNT   int64  = 10                        //IP失败封禁最大失败次数 默认10次
+
 )

innerbean/web_log.go

@@ -98,6 +98,25 @@ func (WebLog) TableName() string {
 	return "web_logs"
 }
 
+// GetIPFailureCount 获取IP在指定时间窗口内的失败次数（用于规则引擎）
+// minutes: 时间窗口（分钟）
+// 返回: 失败次数
+func (w *WebLog) GetIPFailureCount(minutes int64) int64 {
+	if w.SRC_IP == "" {
+		return 0
+	}
+	// 直接调用IP失败管理器（延迟导入避免编译时循环依赖）
+	return getIPFailureCount(w.SRC_IP, minutes)
+}
+
+// getIPFailureCount 获取IP失败次数（通过函数变量实现延迟导入）
+var getIPFailureCount func(string, int64) int64
+
+// SetIPFailureCountGetter 设置IP失败次数获取函数
+func SetIPFailureCountGetter(fn func(string, int64) int64) {
+	getIPFailureCount = fn
+}
+
 type WAFLog struct {
 	REQ_UUID    string `json:"req_uuid"`
 	ACTION      string `json:"action"`

main.go

@@ -15,6 +15,7 @@ import (
 	"SamWaf/wafenginecore"
 	"SamWaf/wafenginecore/wafcaptcha"
 	"SamWaf/wafinit"
+	"SamWaf/wafipban"
 	"SamWaf/wafmangeweb"
 	"SamWaf/wafnotify"
 	"SamWaf/wafowasp"
@@ -210,6 +211,9 @@ func (m *wafSystenService) run() {
 	// 创建owasp
 	global.GWAF_OWASP = wafowasp.NewWafOWASP(true, utils.GetCurrentDir())
 
+	// 初始化ip ban
+	wafipban.InitIPBanManager()
+
 	//提前初始化
 	global.GDATA_CURRENT_LOG_DB_MAP = map[string]*gorm.DB{}
 	rversion := "初始化系统 编译器版本:" + runtime.Version() + " 程序版本号：" + global.GWAF_RELEASE_VERSION_NAME + "(" + global.GWAF_RELEASE_VERSION + ")"

wafipban/ip_failure_manager.go

@@ -0,0 +1,237 @@
+package wafipban
+
+import (
+	"SamWaf/cache"
+	"SamWaf/common/zlog"
+	"SamWaf/enums"
+	"SamWaf/global"
+	"SamWaf/innerbean"
+	"regexp"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+)
+
+func InitIPBanManager() {
+	// 注册到innerbean包，供WebLog使用
+	innerbean.SetIPFailureCountGetter(func(ip string, minutes int64) int64 {
+		return GetIPFailureManager().GetFailureCount(ip, minutes)
+	})
+}
+
+// IPFailureRecord IP失败记录
+type IPFailureRecord struct {
+	IP        string
+	Count     int64
+	FirstTime time.Time
+	LastTime  time.Time
+}
+
+// IPFailureManager IP失败管理器
+type IPFailureManager struct {
+	cache     *cache.WafCache
+	mu        sync.RWMutex
+	statusRe  *regexp.Regexp // 状态码正则表达式
+	statusMap map[int]bool   // 状态码快速查找map
+}
+
+var (
+	ipFailureManagerInstance *IPFailureManager
+	ipFailureManagerOnce     sync.Once
+)
+
+// GetIPFailureManager 获取IP失败管理器单例
+func GetIPFailureManager() *IPFailureManager {
+	ipFailureManagerOnce.Do(func() {
+		ipFailureManagerInstance = &IPFailureManager{
+			cache:     cache.InitWafCache(),
+			statusMap: make(map[int]bool),
+		}
+		ipFailureManagerInstance.initStatusCodes()
+	})
+	return ipFailureManagerInstance
+}
+
+// initStatusCodes 初始化状态码配置
+func (m *IPFailureManager) initStatusCodes() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	statusCodesStr := global.GCONFIG_IP_FAILURE_STATUS_CODES
+	if statusCodesStr == "" {
+		statusCodesStr = "401|403|404|444|429|503"
+	}
+
+	// 清空现有状态码
+	m.statusMap = make(map[int]bool)
+
+	// 尝试解析为数字状态码（用|分隔）
+	parts := strings.Split(statusCodesStr, "|")
+	hasRegex := false
+
+	for _, part := range parts {
+		part = strings.TrimSpace(part)
+		if part == "" {
+			continue
+		}
+
+		// 检查是否是正则表达式（包含特殊字符）
+		if strings.ContainsAny(part, "^$.*+?[]{}()|\\") {
+			hasRegex = true
+			break
+		}
+
+		// 尝试解析为数字
+		if code, err := strconv.Atoi(part); err == nil {
+			m.statusMap[code] = true
+		}
+	}
+
+	// 如果有正则表达式，编译它
+	if hasRegex {
+		re, err := regexp.Compile("^(" + statusCodesStr + ")$")
+		if err != nil {
+			zlog.Warn("IP失败状态码正则表达式编译失败", "error", err.Error(), "pattern", statusCodesStr)
+		} else {
+			m.statusRe = re
+		}
+	}
+}
+
+// ReloadStatusCodes 重新加载状态码配置
+func (m *IPFailureManager) ReloadStatusCodes() {
+	m.initStatusCodes()
+}
+
+// IsFailureStatusCode 检查状态码是否为失败状态码
+func (m *IPFailureManager) IsFailureStatusCode(statusCode int) bool {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	// 先检查快速查找map
+	if m.statusMap[statusCode] {
+		return true
+	}
+
+	// 如果有正则表达式，使用正则匹配
+	if m.statusRe != nil {
+		statusCodeStr := strconv.Itoa(statusCode)
+		return m.statusRe.MatchString(statusCodeStr)
+	}
+
+	return false
+}
+
+// RecordFailure 记录IP失败
+func (m *IPFailureManager) RecordFailure(ip string) {
+	if ip == "" || global.GCONFIG_IP_FAILURE_BAN_ENABLED == 0 {
+		return
+	}
+
+	key := enums.CACHE_IP_FAILURE_PRE + ip
+	now := time.Now()
+
+	// 获取现有记录
+	var record *IPFailureRecord
+	if val := m.cache.Get(key); val != nil {
+		if r, ok := val.(*IPFailureRecord); ok {
+			record = r
+		}
+	}
+
+	// 如果记录不存在或已过期，创建新记录
+	if record == nil {
+		record = &IPFailureRecord{
+			IP:        ip,
+			Count:     1,
+			FirstTime: now,
+			LastTime:  now,
+		}
+	} else {
+		// 检查时间窗口
+		timeWindow := time.Duration(global.GCONFIG_IP_FAILURE_BAN_TIME_WINDOW) * time.Minute
+		if now.Sub(record.FirstTime) > timeWindow {
+			// 超出时间窗口，重置计数
+			record.Count = 1
+			record.FirstTime = now
+			record.LastTime = now
+		} else {
+			// 在时间窗口内，增加计数
+			record.Count++
+			record.LastTime = now
+		}
+	}
+
+	// 保存到缓存，TTL设置为时间窗口的2倍
+	ttl := time.Duration(global.GCONFIG_IP_FAILURE_BAN_TIME_WINDOW*2) * time.Minute
+	m.cache.SetWithTTlRenewTime(key, record, ttl)
+}
+
+// GetFailureCount 获取IP在指定时间窗口内的失败次数
+// minutes: 时间窗口（分钟）
+func (m *IPFailureManager) GetFailureCount(ip string, minutes int64) int64 {
+	if ip == "" || global.GCONFIG_IP_FAILURE_BAN_ENABLED == 0 {
+		return 0
+	}
+
+	key := enums.CACHE_IP_FAILURE_PRE + ip
+	val := m.cache.Get(key)
+	if val == nil {
+		return 0
+	}
+
+	record, ok := val.(*IPFailureRecord)
+	if !ok {
+		return 0
+	}
+
+	// 检查时间窗口
+	timeWindow := time.Duration(minutes) * time.Minute
+	now := time.Now()
+	if now.Sub(record.FirstTime) > timeWindow {
+		// 超出时间窗口，返回0
+		return 0
+	}
+
+	return record.Count
+}
+
+// IsIPBanned 检查IP是否应该被封禁
+func (m *IPFailureManager) IsIPBanned(ip string) bool {
+	if ip == "" || global.GCONFIG_IP_FAILURE_BAN_ENABLED == 0 {
+		return false
+	}
+
+	count := m.GetFailureCount(ip, global.GCONFIG_IP_FAILURE_BAN_TIME_WINDOW)
+	return count >= global.GCONFIG_IP_FAILURE_BAN_MAX_COUNT
+}
+
+// ClearIPFailure 清除IP的失败记录
+func (m *IPFailureManager) ClearIPFailure(ip string) {
+	if ip == "" {
+		return
+	}
+	key := enums.CACHE_IP_FAILURE_PRE + ip
+	m.cache.Remove(key)
+}
+
+// GetFailureInfo 获取IP失败信息（用于调试）
+func (m *IPFailureManager) GetFailureInfo(ip string) *IPFailureRecord {
+	if ip == "" {
+		return nil
+	}
+
+	key := enums.CACHE_IP_FAILURE_PRE + ip
+	val := m.cache.Get(key)
+	if val == nil {
+		return nil
+	}
+
+	record, ok := val.(*IPFailureRecord)
+	if !ok {
+		return nil
+	}
+
+	return record
+}

wafqueue/log_queue.go

@@ -4,6 +4,7 @@ import (
 	"SamWaf/common/zlog"
 	"SamWaf/global"
 	"SamWaf/innerbean"
+	"SamWaf/wafipban"
 	"SamWaf/waftask"
 	"strconv"
 	"time"
@@ -50,6 +51,15 @@ func ProcessLogDequeEngine() {
 				}
 				if len(webLogArray) > 0 {
 					zlog.Debug("日志队列处理协程处理日志数量:" + strconv.Itoa(len(webLogArray)))
+					// 检查失败状态码并记录IP失败
+					if global.GCONFIG_IP_FAILURE_BAN_ENABLED == 1 {
+						ipManager := wafipban.GetIPFailureManager()
+						for _, log := range webLogArray {
+							if ipManager.IsFailureStatusCode(log.STATUS_CODE) {
+								ipManager.RecordFailure(log.SRC_IP)
+							}
+						}
+					}
 					if global.GCONFIG_LOG_PERSIST_ENABLED == 1 {
 						global.GWAF_LOCAL_LOG_DB.CreateInBatches(webLogArray, len(webLogArray))
 					}

waftask/task_config.go

@@ -5,6 +5,7 @@ import (
 	"SamWaf/global"
 	"SamWaf/model"
 	"SamWaf/model/request"
+	"SamWaf/wafipban"
 	"strconv"
 )
 
@@ -113,6 +114,15 @@ func setConfigIntValue(name string, value int64, change int) {
 	case "ip_tag_db":
 		global.GDATA_IP_TAG_DB = value
 		break
+	case "ip_failure_ban_enabled":
+		global.GCONFIG_IP_FAILURE_BAN_ENABLED = value
+		break
+	case "ip_failure_ban_time_window":
+		global.GCONFIG_IP_FAILURE_BAN_TIME_WINDOW = value
+		break
+	case "ip_failure_ban_max_count":
+		global.GCONFIG_IP_FAILURE_BAN_MAX_COUNT = value
+		break
 	default:
 		zlog.Warn("Unknown config item:", name)
 	}
@@ -160,6 +170,11 @@ func setConfigStringValue(name string, value string, change int) {
 	case "ssl_max_version":
 		global.GCONFIG_RECORD_SSLMaxVerson = value
 		break
+	case "ip_failure_status_codes":
+		global.GCONFIG_IP_FAILURE_STATUS_CODES = value
+		// 重新加载状态码配置
+		wafipban.GetIPFailureManager().ReloadStatusCodes()
+		break
 	default:
 		zlog.Warn("Unknown config item:", name)
 	}
@@ -279,4 +294,10 @@ func TaskLoadSetting(initLoad bool) {
 	updateConfigIntItem(initLoad, "database", "batch_insert", global.GDATA_BATCH_INSERT, "数据库批量插入数量", "int", "", configMap)
 	updateConfigIntItem(initLoad, "database", "log_persist_enable", global.GCONFIG_LOG_PERSIST_ENABLED, "是否开启日志持久化（1开启 0关闭）", "options", "0|关闭,1|开启", configMap)
 	updateConfigIntItem(initLoad, "database", "ip_tag_db", global.GDATA_IP_TAG_DB, "IP Tag 存放位置 0 是主库  1是读取 stat库", "int", "", configMap)
+
+	// IP失败封禁相关配置
+	updateConfigStringItem(initLoad, "security", "ip_failure_status_codes", global.GCONFIG_IP_FAILURE_STATUS_CODES, "失败状态码配置，支持多个用|分隔，也支持正则表达式，例如：401|403|404|444|429|503 或 ^4[0-9]{2}$", "string", "", configMap)
+	updateConfigIntItem(initLoad, "security", "ip_failure_ban_enabled", global.GCONFIG_IP_FAILURE_BAN_ENABLED, "是否启用IP失败封禁（1启用 0禁用）", "options", "0|禁用,1|启用", configMap)
+	updateConfigIntItem(initLoad, "security", "ip_failure_ban_time_window", global.GCONFIG_IP_FAILURE_BAN_TIME_WINDOW, "IP失败封禁时间窗口（单位：分钟，默认5分钟）", "int", "", configMap)
+	updateConfigIntItem(initLoad, "security", "ip_failure_ban_max_count", global.GCONFIG_IP_FAILURE_BAN_MAX_COUNT, "IP失败封禁最大失败次数（默认10次）", "int", "", configMap)
 }