SamWaf/waftask/task_batch.go

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"
)

var (
	wafBatchTaskService  = waf_service.WafBatchServiceApp
	wafIPAllowService    = waf_service.WafWhiteIpServiceApp
	wafBlockAllowService = waf_service.WafBlockIpServiceApp
)

/*
*
批量任务
*/
func BatchTask() {
	innerLogName := "BatchTask"
	zlog.Info(innerLogName, "准备进行自动执行批量任务")

	batchTaskList, size, err := wafBatchTaskService.GetAllListInner()
	if err != nil {
		zlog.Error(innerLogName, "批量任务:", err)
		return
	}
	if size <= 0 {
		zlog.Info(innerLogName, "没有需要批量执行的任务")
		return
	}
	for _, batchTask := range batchTaskList {

		switch batchTask.BatchType {
		case enums.BATCHTASK_IPALLOW:
			IPAllowBatch(batchTask)
			break
		case enums.BATCHTASK_IPDENY:
			IPDenyBatch(batchTask)
			break
		}
		zlog.Info(innerLogName, "批量已处理完")

	}
}

// 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
	}
	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
	}

}

// 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
	}
	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
	}

}