feat:add batch sensitive

#276
2026-09-01 15:32:55 +08:00 · 2025-06-04 08:42:27 +08:00
parent 30e4d566af
commit 8816b25c8d
12 changed files with 266 additions and 16 deletions
@@ -106,6 +106,8 @@ func (s *WafBatchTaskApi) ManualBatchTaskApi(c *gin.Context) {
 		case enums.BATCHTASK_IPDENY:
 			waftask.IPDenyBatch(bean)
 			break
+		case enums.BATCHTASK_SENSITIVE:
+			waftask.SensitiveBatch(bean)
 		}
 		response.OkWithMessage("手工执行任务成功", c)
 	} else {
@@ -3,6 +3,7 @@ package enums
 const (
 	BATCHTASK_IPALLOW                = "ipallow"
 	BATCHTASK_IPDENY                 = "ipdeny"
+	BATCHTASK_SENSITIVE              = "sensitive"
 	BATCHTASK_EXECUTEMETHODAPPEND    = "append"
 	BATCHTASK_EXECUTEMETHODOVERWRITE = "overwrite"
 )
@@ -12,5 +12,6 @@ type BatchTask struct {
 	BatchTriggerType   string `json:"batch_trigger_type"`   //触发类型 定时任务 cron ,手动任务 manual
 	BatchSource        string `json:"batch_source"`         //来源内容 路径或者实际的url内容
 	BatchExecuteMethod string `json:"batch_execute_method"` //任务执行方式 追加,覆盖
+	BatchExtraConfig   string `json:"batch_extra_config"`   //额外配置字段(JSON字符串)
 	Remark             string `json:"remark"`               //备注
 }
@@ -10,6 +10,7 @@ type BatchTaskAddReq struct {
 	BatchTriggerType   string `json:"batch_trigger_type"`   //触发类型 定时任务 cron ,手动任务 manual
 	BatchSource        string `json:"batch_source"`         //来源内容 路径或者实际的url内容
 	BatchExecuteMethod string `json:"batch_execute_method"` //任务执行方式 追加,覆盖
+	BatchExtraConfig   string `json:"batch_extra_config"`   //额外配置字段(JSON字符串)
 	Remark             string `json:"remark"`               //备注
 }
 type BatchTaskEditReq struct {
@@ -21,6 +22,7 @@ type BatchTaskEditReq struct {
 	BatchTriggerType   string `json:"batch_trigger_type"`   //触发类型 默认定时任务 cron ,手动任务 manual
 	BatchSource        string `json:"batch_source"`         //来源内容 路径或者实际的url内容
 	BatchExecuteMethod string `json:"batch_execute_method"` //任务执行方式 追加,覆盖
+	BatchExtraConfig   string `json:"batch_extra_config"`   //额外配置字段(JSON字符串)
 	Remark             string `json:"remark"`               //备注
 }
 type BatchTaskDetailReq struct {
@@ -37,6 +37,7 @@ func (receiver *WafBatchTaskService) AddApi(req request.BatchTaskAddReq) error {
 		BatchSourceType:    req.BatchSourceType,
 		BatchTriggerType:   req.BatchTriggerType,
 		BatchType:          req.BatchType,
+		BatchExtraConfig:   req.BatchExtraConfig,
 		Remark:             req.Remark,
 	}
 	global.GWAF_LOCAL_DB.Create(bean)
@@ -63,6 +64,7 @@ func (receiver *WafBatchTaskService) ModifyApi(req request.BatchTaskEditReq) err
 		"BatchSourceType":    req.BatchSourceType,
 		"BatchTriggerType":   req.BatchTriggerType,
 		"BatchType":          req.BatchType,
+		"BatchExtraConfig":   req.BatchExtraConfig,
 		"Remark":             req.Remark,
 	}
 	err := global.GWAF_LOCAL_DB.Model(model.BatchTask{}).Where("id = ?", req.Id).Updates(beanMap).Error
@@ -177,6 +177,14 @@ func pathCoreSql(db *gorm.DB) {
 		zlog.Info("db", "batch_tasks: batch_trigger_type init successfully")
 	}

+	//20250603 批量任务额外配置初始化
+	err = db.Exec("UPDATE batch_tasks SET batch_extra_config='{}' WHERE batch_extra_config IS NULL").Error
+	if err != nil {
+		panic("failed to batch_tasks: batch_extra_config " + err.Error())
+	} else {
+		zlog.Info("db", "batch_tasks: batch_extra_config init successfully")
+	}
+
 	// 记录结束时间并计算耗时
 	duration := time.Since(startTime)
 	zlog.Info("create core default value completely", "duration", duration.String())
@@ -4,7 +4,6 @@ import (
 	"SamWaf/common/zlog"
 	"SamWaf/global"
 	"SamWaf/model"
-	"SamWaf/utils"
 	"bufio"
 	"fmt"
 	"io"
@@ -20,7 +19,7 @@ type BatchProcessor interface {
 	// ProcessItem 处理单个项目
 	ProcessBatch(items []string, task model.BatchTask, progress *BatchProgress) bool
 	// GetExistingItems 获取已存在的项目
-	GetExistingItems(items []string, task model.BatchTask) map[string]interface{}
+	GetExistingItems(items []string, task model.BatchTask, config interface{}) map[string]interface{}
 	// NotifyEngine 通知引擎更新
 	NotifyEngine(task model.BatchTask)
 }
@@ -79,8 +78,11 @@ func ProcessBatchTask(task model.BatchTask, processor BatchProcessor, config Bat
 		return
 	}

+	// 获取对应类型的提取器
+	extractor := GetExtractor(task.BatchType)
+
 	// 首先计算总行数，用于进度显示
-	totalLines, validLines, err := countLines(task)
+	totalLines, validLines, err := countLinesWithExtractor(task, extractor)
 	if err != nil {
 		zlog.Error(innerLogName, fmt.Sprintf("计算总行数失败: %s", err.Error()))
 		// 继续执行，但无法显示准确进度
@@ -107,13 +109,13 @@ func ProcessBatchTask(task model.BatchTask, processor BatchProcessor, config Bat
 			continue // 跳过空行
 		}

-		line = extractIPFromLine(line)                // 可以根据需要替换为更通用的提取函数
-		validRet, _ := utils.IsValidIPOrNetwork(line) // 可以根据需要替换为更通用的验证函数
-		if !validRet {
+		// 使用特定类型的提取器
+		item := extractor.ExtractItem(line)
+		if !extractor.ValidateItem(item) {
 			continue
 		}

-		validItems = append(validItems, line)
+		validItems = append(validItems, item)
 		batchCount++

 		// 当收集到一批或者是最后一批时，进行批量处理
@@ -163,8 +165,8 @@ func ProcessBatchTask(task model.BatchTask, processor BatchProcessor, config Bat
 	}
 }

-// countLines 计算文件总行数和有效行数
-func countLines(task model.BatchTask) (int, int, error) {
+// countLinesWithExtractor 使用提取器计算文件总行数和有效行数
+func countLinesWithExtractor(task model.BatchTask, extractor ItemExtractor) (int, int, error) {
 	contentReader, err := openSource(task)
 	if err != nil {
 		return 0, 0, err
@@ -184,9 +186,8 @@ func countLines(task model.BatchTask) (int, int, error) {
 			continue // 跳过空行
 		}

-		line = extractIPFromLine(line)
-		validRet, _ := utils.IsValidIPOrNetwork(line)
-		if validRet {
+		item := extractor.ExtractItem(line)
+		if extractor.ValidateItem(item) {
 			validLines++
 		}
 	}
@@ -26,7 +26,7 @@ func (p *IPAllowProcessor) ProcessBatch(items []string, task model.BatchTask, pr
 	zlog.Info(logName, fmt.Sprintf("处理白名单批次，包含 %d 个IP", len(items)))

 	// 获取已存在的记录
-	existMap := p.GetExistingItems(items, task)
+	existMap := p.GetExistingItems(items, task, nil)

 	// 根据执行方法处理
 	if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODAPPEND {
@@ -39,7 +39,7 @@ func (p *IPAllowProcessor) ProcessBatch(items []string, task model.BatchTask, pr
 }

 // GetExistingItems 获取已存在的IP记录
-func (p *IPAllowProcessor) GetExistingItems(items []string, task model.BatchTask) map[string]interface{} {
+func (p *IPAllowProcessor) GetExistingItems(items []string, task model.BatchTask, config interface{}) map[string]interface{} {
 	existMap := make(map[string]interface{})
 	var existIPs []model.IPAllowList

@@ -26,7 +26,7 @@ func (p *IPDenyProcessor) ProcessBatch(items []string, task model.BatchTask, pro
 	zlog.Info(logName, fmt.Sprintf("处理黑名单批次，包含 %d 个IP", len(items)))

 	// 获取已存在的记录
-	existMap := p.GetExistingItems(items, task)
+	existMap := p.GetExistingItems(items, task, nil)

 	// 根据执行方法处理
 	if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODAPPEND {
@@ -39,7 +39,7 @@ func (p *IPDenyProcessor) ProcessBatch(items []string, task model.BatchTask, pro
 }

 // GetExistingItems 获取已存在的IP记录
-func (p *IPDenyProcessor) GetExistingItems(items []string, task model.BatchTask) map[string]interface{} {
+func (p *IPDenyProcessor) GetExistingItems(items []string, task model.BatchTask, config interface{}) map[string]interface{} {
 	existMap := make(map[string]interface{})
 	var existIPs []model.IPBlockList

@@ -57,11 +57,26 @@ func (e *DefaultExtractor) ValidateItem(item string) bool {
 	return item != ""
 }

+// SensitiveExtractor 敏感词提取器
+type SensitiveExtractor struct{}
+
+// ExtractItem 敏感词提取，去除前后空格
+func (e *SensitiveExtractor) ExtractItem(line string) string {
+	return strings.TrimSpace(line)
+}
+
+// ValidateItem 敏感词验证，非空且长度合理
+func (e *SensitiveExtractor) ValidateItem(item string) bool {
+	return item != "" && len(item) <= 1000 // 限制敏感词最大长度
+}
+
 // GetExtractor 根据批量任务类型获取合适的提取器
 func GetExtractor(batchType string) ItemExtractor {
 	switch batchType {
 	case "ipallow", "ipdeny":
 		return &IPExtractor{}
+	case "sensitive":
+		return &SensitiveExtractor{}
 	default:
 		return &DefaultExtractor{}
 	}
@@ -0,0 +1,205 @@
+package batch
+
+import (
+	"SamWaf/common/uuid"
+	"SamWaf/common/zlog"
+	"SamWaf/customtype"
+	"SamWaf/enums"
+	"SamWaf/global"
+	"SamWaf/model"
+	"SamWaf/model/baseorm"
+	"SamWaf/model/spec"
+	"encoding/json"
+	"fmt"
+	"time"
+)
+
+// SensitiveConfig 敏感词批量任务额外配置
+type SensitiveConfig struct {
+	CheckDirection string `json:"check_direction"` // 敏感词检测方向 in,out,all
+	Action         string `json:"action"`          // 敏感词检测后动作 deny,replace
+}
+
+// SensitiveProcessor 敏感词处理器
+type SensitiveProcessor struct{}
+
+// ProcessBatch 处理一批敏感词
+func (p *SensitiveProcessor) ProcessBatch(items []string, task model.BatchTask, progress *BatchProgress) bool {
+	if len(items) == 0 {
+		return false
+	}
+
+	logName := "BatchTask-SensitiveBatch"
+	zlog.Info(logName, fmt.Sprintf("处理敏感词批次，包含 %d 个敏感词", len(items)))
+
+	// 解析额外配置
+	var config SensitiveConfig
+	if task.BatchExtraConfig != "" {
+		if err := json.Unmarshal([]byte(task.BatchExtraConfig), &config); err != nil {
+			zlog.Error(logName, "解析敏感词配置失败: "+err.Error())
+			return false
+		}
+	} else {
+		// 默认配置
+		config.CheckDirection = "out"
+		config.Action = "replace"
+	}
+
+	// 获取已存在的记录
+	existMap := p.GetExistingItems(items, task, config)
+
+	// 根据执行方法处理
+	if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODAPPEND {
+		return p.processAppendBatch(items, existMap, task, config, logName, progress)
+	} else if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODOVERWRITE {
+		return p.processOverwriteBatch(items, existMap, task, config, logName, progress)
+	}
+
+	return false
+}
+
+// GetExistingItems 获取已存在的敏感词记录
+func (p *SensitiveProcessor) GetExistingItems(items []string, task model.BatchTask, config interface{}) map[string]interface{} {
+	existMap := make(map[string]interface{})
+	var existSensitives []model.Sensitive
+
+	sensitiveConfig := config.(SensitiveConfig)
+	// 使用IN查询一次性获取所有已存在的记录
+	global.GWAF_LOCAL_DB.Where("content IN (?) and check_direction = ?", items, sensitiveConfig.CheckDirection).Find(&existSensitives)
+
+	// 将已存在的敏感词放入map中，方便快速查找
+	for _, sensitive := range existSensitives {
+		existMap[sensitive.Content] = sensitive
+	}
+
+	return existMap
+}
+
+// NotifyEngine 通知引擎更新
+func (p *SensitiveProcessor) NotifyEngine(task model.BatchTask) {
+	var sensitives []model.Sensitive
+	global.GWAF_LOCAL_DB.Find(&sensitives)
+	var chanInfo = spec.ChanCommonHost{
+		HostCode: task.BatchHostCode,
+		Type:     enums.ChanTypeSensitive,
+		Content:  sensitives,
+	}
+	global.GWAF_CHAN_MSG <- chanInfo
+}
+
+// processAppendBatch 处理追加模式的批次
+func (p *SensitiveProcessor) processAppendBatch(items []string, existMap map[string]interface{}, task model.BatchTask, config SensitiveConfig, logName string, progress *BatchProgress) bool {
+	// 收集需要插入的记录
+	var toInsert []model.Sensitive
+	for _, content := range items {
+		// 如果敏感词不存在，则添加到待插入列表
+		if _, exists := existMap[content]; !exists {
+			toInsert = append(toInsert, model.Sensitive{
+				BaseOrm: baseorm.BaseOrm{
+					Id:          uuid.GenUUID(),
+					USER_CODE:   global.GWAF_USER_CODE,
+					Tenant_ID:   global.GWAF_TENANT_ID,
+					CREATE_TIME: customtype.JsonTime(time.Now()),
+					UPDATE_TIME: customtype.JsonTime(time.Now()),
+				},
+				CheckDirection: config.CheckDirection,
+				Action:         config.Action,
+				Content:        content,
+				Remarks:        time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id,
+			})
+		}
+	}
+
+	// 批量插入新记录
+	if len(toInsert) > 0 {
+		tx := global.GWAF_LOCAL_DB.Begin()
+		if err := tx.Create(&toInsert).Error; err != nil {
+			tx.Rollback()
+			zlog.Error(logName, "批量插入敏感词失败: "+err.Error())
+			return false
+		}
+		tx.Commit()
+
+		zlog.Info(logName, fmt.Sprintf("成功插入 %d 条敏感词记录", len(toInsert)))
+		// 更新进度统计
+		progress.AddInserted(len(toInsert))
+		return len(toInsert) > 0
+	}
+
+	return false
+}
+
+// processOverwriteBatch 处理覆写模式的批次
+func (p *SensitiveProcessor) processOverwriteBatch(items []string, existMap map[string]interface{}, task model.BatchTask, config SensitiveConfig, logName string, progress *BatchProgress) bool {
+	// 收集需要插入和更新的记录
+	var toInsert []model.Sensitive
+	var toUpdate []model.Sensitive
+
+	for _, content := range items {
+		if existSensitive, exists := existMap[content]; exists {
+			// 已存在，需要更新
+			sensitiveRecord := existSensitive.(model.Sensitive)
+			sensitiveRecord.CheckDirection = config.CheckDirection
+			sensitiveRecord.Action = config.Action
+			sensitiveRecord.Remarks = time.Now().Format("20060102") + "批量导入编辑 任务ID:" + task.Id
+			sensitiveRecord.UPDATE_TIME = customtype.JsonTime(time.Now())
+			toUpdate = append(toUpdate, sensitiveRecord)
+		} else {
+			// 不存在，需要插入
+			toInsert = append(toInsert, model.Sensitive{
+				BaseOrm: baseorm.BaseOrm{
+					Id:          uuid.GenUUID(),
+					USER_CODE:   global.GWAF_USER_CODE,
+					Tenant_ID:   global.GWAF_TENANT_ID,
+					CREATE_TIME: customtype.JsonTime(time.Now()),
+					UPDATE_TIME: customtype.JsonTime(time.Now()),
+				},
+				CheckDirection: config.CheckDirection,
+				Action:         config.Action,
+				Content:        content,
+				Remarks:        time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id,
+			})
+		}
+	}
+
+	// 开始事务
+	tx := global.GWAF_LOCAL_DB.Begin()
+	hasChanges := false
+
+	// 批量插入新记录
+	if len(toInsert) > 0 {
+		if err := tx.Create(&toInsert).Error; err != nil {
+			tx.Rollback()
+			zlog.Error(logName, "批量插入敏感词失败: "+err.Error())
+			return false
+		}
+		zlog.Info(logName, fmt.Sprintf("成功插入 %d 条敏感词记录", len(toInsert)))
+		// 更新进度统计
+		progress.AddInserted(len(toInsert))
+		hasChanges = true
+	}
+
+	// 批量更新已存在的记录
+	if len(toUpdate) > 0 {
+		for _, record := range toUpdate {
+			if err := tx.Model(&model.Sensitive{}).Where("id = ?", record.Id).Updates(map[string]interface{}{
+				"CheckDirection": record.CheckDirection,
+				"Action":         record.Action,
+				"Remarks":        record.Remarks,
+				"UPDATE_TIME":    record.UPDATE_TIME,
+			}).Error; err != nil {
+				tx.Rollback()
+				zlog.Error(logName, "批量更新敏感词失败: "+err.Error())
+				return false
+			}
+		}
+		zlog.Info(logName, fmt.Sprintf("成功更新 %d 条敏感词记录", len(toUpdate)))
+		// 更新进度统计
+		progress.AddUpdated(len(toUpdate))
+		hasChanges = true
+	}
+
+	// 提交事务
+	tx.Commit()
+	return hasChanges
+}
@@ -38,6 +38,9 @@ func BatchTask() {
 		case enums.BATCHTASK_IPDENY:
 			IPDenyBatch(batchTask)
 			break
+		case enums.BATCHTASK_SENSITIVE:
+			SensitiveBatch(batchTask)
+			break
 		}
 		zlog.Info(innerLogName, "批量已处理完")

@@ -63,3 +66,13 @@ func IPDenyBatch(task model.BatchTask) {
 	}
 	batch.ProcessBatchTask(task, processor, config)
 }
+
+// SensitiveBatch 敏感词批量处理
+func SensitiveBatch(task model.BatchTask) {
+	processor := &batch.SensitiveProcessor{}
+	config := batch.BatchProcessorConfig{
+		BatchSize: 1000,
+		LogPrefix: "BatchTask-SensitiveBatch",
+	}
+	batch.ProcessBatchTask(task, processor, config)
+}