mirror of
https://gitee.com/samwaf/SamWaf.git
synced 2026-11-01 02:31:39 +08:00
samwaf
@@ -0,0 +1,237 @@
|
||||
package batch
|
||||
|
||||
import (
|
||||
"SamWaf/common/zlog"
|
||||
"SamWaf/global"
|
||||
"SamWaf/model"
|
||||
"SamWaf/utils"
|
||||
"bufio"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// BatchProcessor 批量处理器接口
|
||||
type BatchProcessor interface {
|
||||
// ProcessItem 处理单个项目
|
||||
ProcessBatch(items []string, task model.BatchTask, progress *BatchProgress) bool
|
||||
// GetExistingItems 获取已存在的项目
|
||||
GetExistingItems(items []string, task model.BatchTask) map[string]interface{}
|
||||
// NotifyEngine 通知引擎更新
|
||||
NotifyEngine(task model.BatchTask)
|
||||
}
|
||||
|
||||
// BatchProcessorConfig 批量处理器配置
|
||||
type BatchProcessorConfig struct {
|
||||
BatchSize int // 批处理大小
|
||||
LogPrefix string // 日志前缀
|
||||
}
|
||||
|
||||
// BatchProgress 批处理进度
|
||||
type BatchProgress struct {
|
||||
TotalItems int32 // 总项目数
|
||||
ProcessedItems int32 // 已处理项目数
|
||||
InsertedItems int32 // 已插入项目数
|
||||
UpdatedItems int32 // 已更新项目数
|
||||
}
|
||||
|
||||
// AddProcessed 增加已处理数量
|
||||
func (p *BatchProgress) AddProcessed(count int) {
|
||||
atomic.AddInt32(&p.ProcessedItems, int32(count))
|
||||
}
|
||||
|
||||
// AddInserted 增加已插入数量
|
||||
func (p *BatchProgress) AddInserted(count int) {
|
||||
atomic.AddInt32(&p.InsertedItems, int32(count))
|
||||
}
|
||||
|
||||
// AddUpdated 增加已更新数量
|
||||
func (p *BatchProgress) AddUpdated(count int) {
|
||||
atomic.AddInt32(&p.UpdatedItems, int32(count))
|
||||
}
|
||||
|
||||
// GetProgress 获取进度百分比
|
||||
func (p *BatchProgress) GetProgress() float64 {
|
||||
if p.TotalItems == 0 {
|
||||
return 100.0
|
||||
}
|
||||
return float64(p.ProcessedItems) / float64(p.TotalItems) * 100.0
|
||||
}
|
||||
|
||||
// ProcessBatchTask 通用批量处理函数
|
||||
func ProcessBatchTask(task model.BatchTask, processor BatchProcessor, config BatchProcessorConfig) {
|
||||
innerLogName := config.LogPrefix
|
||||
|
||||
contentReader, err := openSource(task)
|
||||
if err != nil {
|
||||
zlog.Error(innerLogName, err.Error())
|
||||
return
|
||||
}
|
||||
defer contentReader.Close()
|
||||
|
||||
// 判断数据库是否已经关闭
|
||||
if global.GWAF_LOCAL_DB == nil {
|
||||
zlog.Error(innerLogName, "数据库已经关闭批量处理终止")
|
||||
return
|
||||
}
|
||||
|
||||
// 首先计算总行数，用于进度显示
|
||||
totalLines, validLines, err := countLines(task)
|
||||
if err != nil {
|
||||
zlog.Error(innerLogName, fmt.Sprintf("计算总行数失败: %s", err.Error()))
|
||||
// 继续执行，但无法显示准确进度
|
||||
}
|
||||
|
||||
// 创建进度跟踪器
|
||||
progress := &BatchProgress{
|
||||
TotalItems: int32(validLines),
|
||||
}
|
||||
|
||||
zlog.Info(innerLogName, fmt.Sprintf("开始批量处理，总行数: %d，有效行数: %d", totalLines, validLines))
|
||||
|
||||
// 收集有效的项目
|
||||
validItems := make([]string, 0, config.BatchSize)
|
||||
scanner := bufio.NewScanner(contentReader)
|
||||
hasAffectInfo := false
|
||||
|
||||
// 按批次处理
|
||||
batchCount := 0
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" {
|
||||
continue // 跳过空行
|
||||
}
|
||||
|
||||
line = extractIPFromLine(line) // 可以根据需要替换为更通用的提取函数
|
||||
validRet, _ := utils.IsValidIPOrNetwork(line) // 可以根据需要替换为更通用的验证函数
|
||||
if !validRet {
|
||||
continue
|
||||
}
|
||||
|
||||
validItems = append(validItems, line)
|
||||
batchCount++
|
||||
|
||||
// 当收集到一批或者是最后一批时，进行批量处理
|
||||
if batchCount >= config.BatchSize {
|
||||
// 处理当前批次
|
||||
if processor.ProcessBatch(validItems, task, progress) {
|
||||
hasAffectInfo = true
|
||||
}
|
||||
|
||||
// 记录进度
|
||||
progress.AddProcessed(len(validItems))
|
||||
zlog.Info(innerLogName, fmt.Sprintf("进度: %.2f%% (%d/%d)，已插入: %d，已更新: %d",
|
||||
progress.GetProgress(), progress.ProcessedItems, progress.TotalItems,
|
||||
progress.InsertedItems, progress.UpdatedItems))
|
||||
|
||||
// 重置批次
|
||||
validItems = make([]string, 0, config.BatchSize)
|
||||
batchCount = 0
|
||||
}
|
||||
}
|
||||
|
||||
// 处理最后一批不足BatchSize的项目
|
||||
if len(validItems) > 0 {
|
||||
if processor.ProcessBatch(validItems, task, progress) {
|
||||
hasAffectInfo = true
|
||||
}
|
||||
|
||||
// 记录进度
|
||||
progress.AddProcessed(len(validItems))
|
||||
zlog.Info(innerLogName, fmt.Sprintf("进度: %.2f%% (%d/%d)，已插入: %d，已更新: %d",
|
||||
progress.GetProgress(), progress.ProcessedItems, progress.TotalItems,
|
||||
progress.InsertedItems, progress.UpdatedItems))
|
||||
}
|
||||
|
||||
if hasAffectInfo {
|
||||
// 通知引擎进行实时生效
|
||||
processor.NotifyEngine(task)
|
||||
}
|
||||
|
||||
// 输出最终处理结果
|
||||
zlog.Info(innerLogName, fmt.Sprintf("批量处理完成，总计处理: %d，插入: %d，更新: %d",
|
||||
progress.ProcessedItems, progress.InsertedItems, progress.UpdatedItems))
|
||||
|
||||
// 检查扫描错误
|
||||
if err := scanner.Err(); err != nil {
|
||||
zlog.Error(innerLogName, fmt.Sprintf("扫描文件时发生错误: %s", err.Error()))
|
||||
}
|
||||
}
|
||||
|
||||
// countLines 计算文件总行数和有效行数
|
||||
func countLines(task model.BatchTask) (int, int, error) {
|
||||
contentReader, err := openSource(task)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
defer contentReader.Close()
|
||||
|
||||
scanner := bufio.NewScanner(contentReader)
|
||||
totalLines := 0
|
||||
validLines := 0
|
||||
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
totalLines++
|
||||
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" {
|
||||
continue // 跳过空行
|
||||
}
|
||||
|
||||
line = extractIPFromLine(line)
|
||||
validRet, _ := utils.IsValidIPOrNetwork(line)
|
||||
if validRet {
|
||||
validLines++
|
||||
}
|
||||
}
|
||||
|
||||
return totalLines, validLines, scanner.Err()
|
||||
}
|
||||
|
||||
// openSource 打开本地或远程数据源
|
||||
func openSource(task model.BatchTask) (io.ReadCloser, error) {
|
||||
if task.BatchSourceType == "local" {
|
||||
f, err := os.Open(task.BatchSource)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to open local file: %v", err)
|
||||
}
|
||||
return f, nil
|
||||
}
|
||||
// remote
|
||||
resp, err := http.Get(task.BatchSource)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to fetch remote data: %v", err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
resp.Body.Close()
|
||||
return nil, fmt.Errorf("unexpected status code: %d", resp.StatusCode)
|
||||
}
|
||||
return resp.Body, nil
|
||||
}
|
||||
|
||||
// extractIPFromLine 使用正则表达式从行中提取IP地址或网段
|
||||
func extractIPFromLine(line string) string {
|
||||
// 匹配IPv4地址或IPv4网段
|
||||
ipv4Regex := regexp.MustCompile(`\b(?:\d{1,3}\.){3}\d{1,3}(?:/\d{1,2})?\b`)
|
||||
|
||||
// 匹配IPv6地址或IPv6网段 (简化版本)
|
||||
ipv6Regex := regexp.MustCompile(`\b(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}(?:/\d{1,3})?\b`)
|
||||
|
||||
// 先尝试匹配IPv4
|
||||
if match := ipv4Regex.FindString(line); match != "" {
|
||||
return match
|
||||
}
|
||||
|
||||
// 再尝试匹配IPv6
|
||||
if match := ipv6Regex.FindString(line); match != "" {
|
||||
return match
|
||||
}
|
||||
|
||||
return line // 如果没有匹配到，返回原始行
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
package batch
|
||||
|
||||
import (
|
||||
"SamWaf/common/uuid"
|
||||
"SamWaf/common/zlog"
|
||||
"SamWaf/customtype"
|
||||
"SamWaf/enums"
|
||||
"SamWaf/global"
|
||||
"SamWaf/model"
|
||||
"SamWaf/model/baseorm"
|
||||
"SamWaf/model/spec"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// IPAllowProcessor 白名单IP处理器
|
||||
type IPAllowProcessor struct{}
|
||||
|
||||
// ProcessBatch 处理一批IP
|
||||
func (p *IPAllowProcessor) ProcessBatch(items []string, task model.BatchTask, progress *BatchProgress) bool {
|
||||
if len(items) == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
logName := "BatchTask-IPAllowBatch"
|
||||
zlog.Info(logName, fmt.Sprintf("处理白名单批次，包含 %d 个IP", len(items)))
|
||||
|
||||
// 获取已存在的记录
|
||||
existMap := p.GetExistingItems(items, task)
|
||||
|
||||
// 根据执行方法处理
|
||||
if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODAPPEND {
|
||||
return p.processAppendBatch(items, existMap, task, logName, progress)
|
||||
} else if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODOVERWRITE {
|
||||
return p.processOverwriteBatch(items, existMap, task, logName, progress)
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// GetExistingItems 获取已存在的IP记录
|
||||
func (p *IPAllowProcessor) GetExistingItems(items []string, task model.BatchTask) map[string]interface{} {
|
||||
existMap := make(map[string]interface{})
|
||||
var existIPs []model.IPAllowList
|
||||
|
||||
// 使用IN查询一次性获取所有已存在的记录
|
||||
global.GWAF_LOCAL_DB.Where("host_code = ? AND ip IN (?)", task.BatchHostCode, items).Find(&existIPs)
|
||||
|
||||
// 将已存在的IP放入map中，方便快速查找
|
||||
for _, ip := range existIPs {
|
||||
existMap[ip.Ip] = ip
|
||||
}
|
||||
|
||||
return existMap
|
||||
}
|
||||
|
||||
// NotifyEngine 通知引擎更新
|
||||
func (p *IPAllowProcessor) NotifyEngine(task model.BatchTask) {
|
||||
var ipWhites []model.IPAllowList
|
||||
global.GWAF_LOCAL_DB.Where("host_code = ? ", task.BatchHostCode).Find(&ipWhites)
|
||||
var chanInfo = spec.ChanCommonHost{
|
||||
HostCode: task.BatchHostCode,
|
||||
Type: enums.ChanTypeAllowIP,
|
||||
Content: ipWhites,
|
||||
}
|
||||
global.GWAF_CHAN_MSG <- chanInfo
|
||||
}
|
||||
|
||||
// processAppendBatch 处理追加模式的批次
|
||||
func (p *IPAllowProcessor) processAppendBatch(items []string, existMap map[string]interface{}, task model.BatchTask, logName string, progress *BatchProgress) bool {
|
||||
// 收集需要插入的记录
|
||||
var toInsert []model.IPAllowList
|
||||
for _, ip := range items {
|
||||
// 如果IP不存在，则添加到待插入列表
|
||||
if _, exists := existMap[ip]; !exists {
|
||||
toInsert = append(toInsert, model.IPAllowList{
|
||||
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()),
|
||||
},
|
||||
HostCode: task.BatchHostCode,
|
||||
Ip: ip,
|
||||
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, "批量插入白名单IP失败: "+err.Error())
|
||||
return false
|
||||
}
|
||||
tx.Commit()
|
||||
|
||||
zlog.Info(logName, fmt.Sprintf("成功插入 %d 条白名单IP记录", len(toInsert)))
|
||||
// 更新进度统计
|
||||
progress.AddInserted(len(toInsert))
|
||||
return len(toInsert) > 0
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// processOverwriteBatch 处理覆写模式的批次
|
||||
func (p *IPAllowProcessor) processOverwriteBatch(items []string, existMap map[string]interface{}, task model.BatchTask, logName string, progress *BatchProgress) bool {
|
||||
// 收集需要插入和更新的记录
|
||||
var toInsert []model.IPAllowList
|
||||
var toUpdate []model.IPAllowList
|
||||
|
||||
for _, ip := range items {
|
||||
if existIP, exists := existMap[ip]; exists {
|
||||
// 已存在，需要更新
|
||||
ipRecord := existIP.(model.IPAllowList)
|
||||
ipRecord.Remarks = time.Now().Format("20060102") + "批量导入编辑 任务ID:" + task.Id
|
||||
ipRecord.UPDATE_TIME = customtype.JsonTime(time.Now())
|
||||
toUpdate = append(toUpdate, ipRecord)
|
||||
} else {
|
||||
// 不存在，需要插入
|
||||
toInsert = append(toInsert, model.IPAllowList{
|
||||
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()),
|
||||
},
|
||||
HostCode: task.BatchHostCode,
|
||||
Ip: ip,
|
||||
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, "批量插入白名单IP失败: "+err.Error())
|
||||
return false
|
||||
}
|
||||
zlog.Info(logName, fmt.Sprintf("成功插入 %d 条白名单IP记录", len(toInsert)))
|
||||
// 更新进度统计
|
||||
progress.AddInserted(len(toInsert))
|
||||
hasChanges = true
|
||||
}
|
||||
|
||||
// 批量更新已存在的记录
|
||||
if len(toUpdate) > 0 {
|
||||
for _, record := range toUpdate {
|
||||
if err := tx.Model(&model.IPAllowList{}).Where("id = ?", record.Id).Updates(map[string]interface{}{
|
||||
"Remarks": record.Remarks,
|
||||
"UPDATE_TIME": record.UPDATE_TIME,
|
||||
}).Error; err != nil {
|
||||
tx.Rollback()
|
||||
zlog.Error(logName, "批量更新白名单IP失败: "+err.Error())
|
||||
return false
|
||||
}
|
||||
}
|
||||
zlog.Info(logName, fmt.Sprintf("成功更新 %d 条白名单IP记录", len(toUpdate)))
|
||||
// 更新进度统计
|
||||
progress.AddUpdated(len(toUpdate))
|
||||
hasChanges = true
|
||||
}
|
||||
|
||||
// 提交事务
|
||||
tx.Commit()
|
||||
return hasChanges
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
package batch
|
||||
|
||||
import (
|
||||
"SamWaf/common/uuid"
|
||||
"SamWaf/common/zlog"
|
||||
"SamWaf/customtype"
|
||||
"SamWaf/enums"
|
||||
"SamWaf/global"
|
||||
"SamWaf/model"
|
||||
"SamWaf/model/baseorm"
|
||||
"SamWaf/model/spec"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// IPDenyProcessor 黑名单IP处理器
|
||||
type IPDenyProcessor struct{}
|
||||
|
||||
// ProcessBatch 处理一批IP
|
||||
func (p *IPDenyProcessor) ProcessBatch(items []string, task model.BatchTask, progress *BatchProgress) bool {
|
||||
if len(items) == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
logName := "BatchTask-IPDenyBatch"
|
||||
zlog.Info(logName, fmt.Sprintf("处理黑名单批次，包含 %d 个IP", len(items)))
|
||||
|
||||
// 获取已存在的记录
|
||||
existMap := p.GetExistingItems(items, task)
|
||||
|
||||
// 根据执行方法处理
|
||||
if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODAPPEND {
|
||||
return p.processAppendBatch(items, existMap, task, logName, progress)
|
||||
} else if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODOVERWRITE {
|
||||
return p.processOverwriteBatch(items, existMap, task, logName, progress)
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// GetExistingItems 获取已存在的IP记录
|
||||
func (p *IPDenyProcessor) GetExistingItems(items []string, task model.BatchTask) map[string]interface{} {
|
||||
existMap := make(map[string]interface{})
|
||||
var existIPs []model.IPBlockList
|
||||
|
||||
// 使用IN查询一次性获取所有已存在的记录
|
||||
global.GWAF_LOCAL_DB.Where("host_code = ? AND ip IN (?)", task.BatchHostCode, items).Find(&existIPs)
|
||||
|
||||
// 将已存在的IP放入map中，方便快速查找
|
||||
for _, ip := range existIPs {
|
||||
existMap[ip.Ip] = ip
|
||||
}
|
||||
|
||||
return existMap
|
||||
}
|
||||
|
||||
// NotifyEngine 通知引擎更新
|
||||
func (p *IPDenyProcessor) NotifyEngine(task model.BatchTask) {
|
||||
var ipBlocks []model.IPBlockList
|
||||
global.GWAF_LOCAL_DB.Where("host_code = ? ", task.BatchHostCode).Find(&ipBlocks)
|
||||
var chanInfo = spec.ChanCommonHost{
|
||||
HostCode: task.BatchHostCode,
|
||||
Type: enums.ChanTypeBlockIP,
|
||||
Content: ipBlocks,
|
||||
}
|
||||
global.GWAF_CHAN_MSG <- chanInfo
|
||||
}
|
||||
|
||||
// processAppendBatch 处理追加模式的批次
|
||||
func (p *IPDenyProcessor) processAppendBatch(items []string, existMap map[string]interface{}, task model.BatchTask, logName string, progress *BatchProgress) bool {
|
||||
// 收集需要插入的记录
|
||||
var toInsert []model.IPBlockList
|
||||
for _, ip := range items {
|
||||
// 如果IP不存在，则添加到待插入列表
|
||||
if _, exists := existMap[ip]; !exists {
|
||||
toInsert = append(toInsert, model.IPBlockList{
|
||||
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()),
|
||||
},
|
||||
HostCode: task.BatchHostCode,
|
||||
Ip: ip,
|
||||
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, "批量插入黑名单IP失败: "+err.Error())
|
||||
return false
|
||||
}
|
||||
tx.Commit()
|
||||
|
||||
zlog.Info(logName, fmt.Sprintf("成功插入 %d 条黑名单IP记录", len(toInsert)))
|
||||
// 更新进度统计
|
||||
progress.AddInserted(len(toInsert))
|
||||
return len(toInsert) > 0
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// processOverwriteBatch 处理覆写模式的批次
|
||||
func (p *IPDenyProcessor) processOverwriteBatch(items []string, existMap map[string]interface{}, task model.BatchTask, logName string, progress *BatchProgress) bool {
|
||||
// 收集需要插入和更新的记录
|
||||
var toInsert []model.IPBlockList
|
||||
var toUpdate []model.IPBlockList
|
||||
|
||||
for _, ip := range items {
|
||||
if existIP, exists := existMap[ip]; exists {
|
||||
// 已存在，需要更新
|
||||
ipRecord := existIP.(model.IPBlockList)
|
||||
ipRecord.Remarks = time.Now().Format("20060102") + "批量导入编辑 任务ID:" + task.Id
|
||||
ipRecord.UPDATE_TIME = customtype.JsonTime(time.Now())
|
||||
toUpdate = append(toUpdate, ipRecord)
|
||||
} else {
|
||||
// 不存在，需要插入
|
||||
toInsert = append(toInsert, model.IPBlockList{
|
||||
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()),
|
||||
},
|
||||
HostCode: task.BatchHostCode,
|
||||
Ip: ip,
|
||||
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, "批量插入黑名单IP失败: "+err.Error())
|
||||
return false
|
||||
}
|
||||
zlog.Info(logName, fmt.Sprintf("成功插入 %d 条黑名单IP记录", len(toInsert)))
|
||||
// 更新进度统计
|
||||
progress.AddInserted(len(toInsert))
|
||||
hasChanges = true
|
||||
}
|
||||
|
||||
// 批量更新已存在的记录
|
||||
if len(toUpdate) > 0 {
|
||||
for _, record := range toUpdate {
|
||||
if err := tx.Model(&model.IPBlockList{}).Where("id = ?", record.Id).Updates(map[string]interface{}{
|
||||
"Remarks": record.Remarks,
|
||||
"UPDATE_TIME": record.UPDATE_TIME,
|
||||
}).Error; err != nil {
|
||||
tx.Rollback()
|
||||
zlog.Error(logName, "批量更新黑名单IP失败: "+err.Error())
|
||||
return false
|
||||
}
|
||||
}
|
||||
zlog.Info(logName, fmt.Sprintf("成功更新 %d 条黑名单IP记录", len(toUpdate)))
|
||||
// 更新进度统计
|
||||
progress.AddUpdated(len(toUpdate))
|
||||
hasChanges = true
|
||||
}
|
||||
|
||||
// 提交事务
|
||||
tx.Commit()
|
||||
return hasChanges
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package batch
|
||||
|
||||
import (
|
||||
"SamWaf/utils"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ItemExtractor 项目提取器接口
|
||||
type ItemExtractor interface {
|
||||
// ExtractItem 从行中提取项目
|
||||
ExtractItem(line string) string
|
||||
// ValidateItem 验证项目是否有效
|
||||
ValidateItem(item string) bool
|
||||
}
|
||||
|
||||
// IPExtractor IP提取器
|
||||
type IPExtractor struct{}
|
||||
|
||||
// ExtractItem 从行中提取IP地址
|
||||
func (e *IPExtractor) ExtractItem(line string) string {
|
||||
// 匹配IPv4地址或IPv4网段
|
||||
ipv4Regex := regexp.MustCompile(`\b(?:\d{1,3}\.){3}\d{1,3}(?:/\d{1,2})?\b`)
|
||||
|
||||
// 匹配IPv6地址或IPv6网段 (简化版本)
|
||||
ipv6Regex := regexp.MustCompile(`\b(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}(?:/\d{1,3})?\b`)
|
||||
|
||||
// 先尝试匹配IPv4
|
||||
if match := ipv4Regex.FindString(line); match != "" {
|
||||
return match
|
||||
}
|
||||
|
||||
// 再尝试匹配IPv6
|
||||
if match := ipv6Regex.FindString(line); match != "" {
|
||||
return match
|
||||
}
|
||||
|
||||
return line // 如果没有匹配到，返回原始行
|
||||
}
|
||||
|
||||
// ValidateItem 验证IP地址是否有效
|
||||
func (e *IPExtractor) ValidateItem(item string) bool {
|
||||
validRet, _ := utils.IsValidIPOrNetwork(item)
|
||||
return validRet
|
||||
}
|
||||
|
||||
// DefaultExtractor 默认提取器，不做特殊处理
|
||||
type DefaultExtractor struct{}
|
||||
|
||||
// ExtractItem 默认提取，只做简单的空格处理
|
||||
func (e *DefaultExtractor) ExtractItem(line string) string {
|
||||
return strings.TrimSpace(line)
|
||||
}
|
||||
|
||||
// ValidateItem 默认验证，非空即有效
|
||||
func (e *DefaultExtractor) ValidateItem(item string) bool {
|
||||
return item != ""
|
||||
}
|
||||
|
||||
// GetExtractor 根据批量任务类型获取合适的提取器
|
||||
func GetExtractor(batchType string) ItemExtractor {
|
||||
switch batchType {
|
||||
case "ipallow", "ipdeny":
|
||||
return &IPExtractor{}
|
||||
default:
|
||||
return &DefaultExtractor{}
|
||||
}
|
||||
}
|
||||
+12
-174
@@ -3,25 +3,13 @@ package waftask
|
||||
import (
|
||||
"SamWaf/common/zlog"
|
||||
"SamWaf/enums"
|
||||
"SamWaf/global"
|
||||
"SamWaf/model"
|
||||
"SamWaf/model/request"
|
||||
"SamWaf/model/spec"
|
||||
"SamWaf/service/waf_service"
|
||||
"SamWaf/utils"
|
||||
"errors"
|
||||
"fmt"
|
||||
"gorm.io/gorm"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
"SamWaf/waftask/batch"
|
||||
)
|
||||
|
||||
var (
|
||||
wafBatchTaskService = waf_service.WafBatchServiceApp
|
||||
wafIPAllowService = waf_service.WafWhiteIpServiceApp
|
||||
wafBlockAllowService = waf_service.WafBlockIpServiceApp
|
||||
wafBatchTaskService = waf_service.WafBatchServiceApp
|
||||
)
|
||||
|
||||
/*
|
||||
@@ -56,172 +44,22 @@ func BatchTask() {
|
||||
}
|
||||
}
|
||||
|
||||
// handleLocalSource 从本地路径读取数据
|
||||
func handleLocalSource(task model.BatchTask) (string, error) {
|
||||
data, err := ioutil.ReadFile(task.BatchSource)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to read local file: %v", err)
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
// handleRemoteSource 从远程 URL 获取数据
|
||||
func handleRemoteSource(task model.BatchTask) (string, error) {
|
||||
resp, err := http.Get(task.BatchSource)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to fetch remote data: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("unexpected status code: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
data, err := ioutil.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to read remote response body: %v", err)
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
// IPAllowBatch 白名单IP批量处理
|
||||
func IPAllowBatch(task model.BatchTask) {
|
||||
innerLogName := "BatchTask-IPAllowBatch"
|
||||
//1.取数据
|
||||
//2.处理数据
|
||||
content := ""
|
||||
if task.BatchSourceType == "local" {
|
||||
cnt, err := handleLocalSource(task)
|
||||
if err != nil {
|
||||
zlog.Error(innerLogName, err.Error())
|
||||
return
|
||||
}
|
||||
content = cnt
|
||||
} else if task.BatchSourceType == "remote" {
|
||||
cnt, err := handleRemoteSource(task)
|
||||
if err != nil {
|
||||
zlog.Error(innerLogName, err.Error())
|
||||
return
|
||||
}
|
||||
content = cnt
|
||||
processor := &batch.IPAllowProcessor{}
|
||||
config := batch.BatchProcessorConfig{
|
||||
BatchSize: 1000,
|
||||
LogPrefix: "BatchTask-IPAllowBatch",
|
||||
}
|
||||
if content == "" {
|
||||
zlog.Error(innerLogName, task.BatchTaskName+"没有数据需要处理")
|
||||
return
|
||||
}
|
||||
hasAffectInfo := false
|
||||
lines := strings.Split(content, "\n")
|
||||
// 遍历每一行数据进行处理
|
||||
for _, line := range lines {
|
||||
if line == "" {
|
||||
continue // 跳过空行
|
||||
}
|
||||
line = strings.TrimSpace(line)
|
||||
validRet, _ := utils.IsValidIPOrNetwork(line)
|
||||
if !validRet {
|
||||
continue
|
||||
}
|
||||
if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODAPPEND {
|
||||
//添加
|
||||
err := wafIPAllowService.CheckIsExistApi(request.WafAllowIpAddReq{task.BatchHostCode, line, ""})
|
||||
if err != nil && errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
wafIPAllowService.AddApi(request.WafAllowIpAddReq{task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id})
|
||||
hasAffectInfo = true
|
||||
}
|
||||
} else if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODOVERWRITE {
|
||||
//覆写
|
||||
bean := wafIPAllowService.GetDetailByIPApi(line, task.BatchHostCode)
|
||||
if bean.HostCode == "" {
|
||||
wafIPAllowService.AddApi(request.WafAllowIpAddReq{task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id})
|
||||
hasAffectInfo = true
|
||||
} else {
|
||||
wafIPAllowService.ModifyApi(request.WafAllowIpEditReq{bean.Id, task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入编辑 任务ID:" + task.Id})
|
||||
hasAffectInfo = true
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
if hasAffectInfo {
|
||||
//发送通知到引擎进行实时生效
|
||||
var ipWhites []model.IPAllowList
|
||||
global.GWAF_LOCAL_DB.Where("host_code = ? ", task.BatchHostCode).Find(&ipWhites)
|
||||
var chanInfo = spec.ChanCommonHost{
|
||||
HostCode: task.BatchHostCode,
|
||||
Type: enums.ChanTypeAllowIP,
|
||||
Content: ipWhites,
|
||||
}
|
||||
global.GWAF_CHAN_MSG <- chanInfo
|
||||
}
|
||||
|
||||
batch.ProcessBatchTask(task, processor, config)
|
||||
}
|
||||
|
||||
// IPDenyBatch 黑名单IP批量处理
|
||||
func IPDenyBatch(task model.BatchTask) {
|
||||
innerLogName := "BatchTask-IPDenyBatch"
|
||||
//1.取数据
|
||||
//2.处理数据
|
||||
content := ""
|
||||
if task.BatchSourceType == "local" {
|
||||
cnt, err := handleLocalSource(task)
|
||||
if err != nil {
|
||||
zlog.Error(innerLogName, err.Error())
|
||||
return
|
||||
}
|
||||
content = cnt
|
||||
} else if task.BatchSourceType == "remote" {
|
||||
cnt, err := handleRemoteSource(task)
|
||||
if err != nil {
|
||||
zlog.Error(innerLogName, err.Error())
|
||||
return
|
||||
}
|
||||
content = cnt
|
||||
processor := &batch.IPDenyProcessor{}
|
||||
config := batch.BatchProcessorConfig{
|
||||
BatchSize: 1000,
|
||||
LogPrefix: "BatchTask-IPDenyBatch",
|
||||
}
|
||||
if content == "" {
|
||||
zlog.Error(innerLogName, task.BatchTaskName+"没有数据需要处理")
|
||||
return
|
||||
}
|
||||
hasAffectInfo := false
|
||||
lines := strings.Split(content, "\n")
|
||||
// 遍历每一行数据进行处理
|
||||
for _, line := range lines {
|
||||
if line == "" {
|
||||
continue // 跳过空行
|
||||
}
|
||||
line = strings.TrimSpace(line)
|
||||
validRet, _ := utils.IsValidIPOrNetwork(line)
|
||||
if !validRet {
|
||||
continue
|
||||
}
|
||||
if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODAPPEND {
|
||||
//添加
|
||||
err := wafBlockAllowService.CheckIsExistApi(request.WafBlockIpAddReq{task.BatchHostCode, line, ""})
|
||||
if err != nil && errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
wafBlockAllowService.AddApi(request.WafBlockIpAddReq{task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id})
|
||||
hasAffectInfo = true
|
||||
}
|
||||
} else if task.BatchExecuteMethod == enums.BATCHTASK_EXECUTEMETHODOVERWRITE {
|
||||
//覆写
|
||||
bean := wafBlockAllowService.GetDetailByIPApi(line, task.BatchHostCode)
|
||||
if bean.HostCode == "" {
|
||||
wafBlockAllowService.AddApi(request.WafBlockIpAddReq{task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入 任务ID:" + task.Id})
|
||||
hasAffectInfo = true
|
||||
} else {
|
||||
wafBlockAllowService.ModifyApi(request.WafBlockIpEditReq{bean.Id, task.BatchHostCode, line, time.Now().Format("20060102") + "批量导入编辑 任务ID:" + task.Id})
|
||||
hasAffectInfo = true
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
if hasAffectInfo {
|
||||
//发送通知到引擎进行实时生效
|
||||
var ipBlocks []model.IPBlockList
|
||||
global.GWAF_LOCAL_DB.Where("host_code = ? ", task.BatchHostCode).Find(&ipBlocks)
|
||||
var chanInfo = spec.ChanCommonHost{
|
||||
HostCode: task.BatchHostCode,
|
||||
Type: enums.ChanTypeBlockIP,
|
||||
Content: ipBlocks,
|
||||
}
|
||||
global.GWAF_CHAN_MSG <- chanInfo
|
||||
}
|
||||
|
||||
batch.ProcessBatchTask(task, processor, config)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user