diff --git a/api/waf_batchtask.go b/api/waf_batchtask.go index cfb6efc..27b3101 100644 --- a/api/waf_batchtask.go +++ b/api/waf_batchtask.go @@ -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 { diff --git a/enums/batch_task_enum.go b/enums/batch_task_enum.go index a7cdf1f..018729a 100644 --- a/enums/batch_task_enum.go +++ b/enums/batch_task_enum.go @@ -3,6 +3,7 @@ package enums const ( BATCHTASK_IPALLOW = "ipallow" BATCHTASK_IPDENY = "ipdeny" + BATCHTASK_SENSITIVE = "sensitive" BATCHTASK_EXECUTEMETHODAPPEND = "append" BATCHTASK_EXECUTEMETHODOVERWRITE = "overwrite" ) diff --git a/model/batchtask.go b/model/batchtask.go index 8de7cb5..7b8fdb5 100644 --- a/model/batchtask.go +++ b/model/batchtask.go @@ -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"` //备注 } diff --git a/model/request/waf_batchtask_req.go b/model/request/waf_batchtask_req.go index 4694cb5..7abd4ad 100644 --- a/model/request/waf_batchtask_req.go +++ b/model/request/waf_batchtask_req.go @@ -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 { diff --git a/service/waf_service/waf_batchtask.go b/service/waf_service/waf_batchtask.go index c713658..a9a10c3 100644 --- a/service/waf_service/waf_batchtask.go +++ b/service/waf_service/waf_batchtask.go @@ -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 diff --git a/wafdb/patch_sql.go b/wafdb/patch_sql.go index 964399b..c9ac4fb 100644 --- a/wafdb/patch_sql.go +++ b/wafdb/patch_sql.go @@ -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()) diff --git a/waftask/batch/batch_processor.go b/waftask/batch/batch_processor.go index dff1673..b8724c2 100644 --- a/waftask/batch/batch_processor.go +++ b/waftask/batch/batch_processor.go @@ -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++ } } diff --git a/waftask/batch/ip_allow_processor.go b/waftask/batch/ip_allow_processor.go index cd3f005..de92ea9 100644 --- a/waftask/batch/ip_allow_processor.go +++ b/waftask/batch/ip_allow_processor.go @@ -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 diff --git a/waftask/batch/ip_deny_processor.go b/waftask/batch/ip_deny_processor.go index 49aba7c..eb7cd41 100644 --- a/waftask/batch/ip_deny_processor.go +++ b/waftask/batch/ip_deny_processor.go @@ -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 diff --git a/waftask/batch/item_extractor.go b/waftask/batch/item_extractor.go index 719ff3e..b6951f2 100644 --- a/waftask/batch/item_extractor.go +++ b/waftask/batch/item_extractor.go @@ -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{} } diff --git a/waftask/batch/sensitive_processor.go b/waftask/batch/sensitive_processor.go new file mode 100644 index 0000000..48c1299 --- /dev/null +++ b/waftask/batch/sensitive_processor.go @@ -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 +} diff --git a/waftask/task_batch.go b/waftask/task_batch.go index 2425afc..009de67 100644 --- a/waftask/task_batch.go +++ b/waftask/task_batch.go @@ -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) +}