opensource/SamWaf
3
0
Fork 0
mirror of https://gitee.com/samwaf/SamWaf.git synced 2026-09-01 15:32:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #828 from samwafgo/feat_ai

Browse Source 
feat:add light ai model
This commit is contained in:
samwafgo
2026-06-15 17:04:19 +08:00
committed by GitHub
parent 17805e19ab 17ac5ebbd3
commit 9366e0962b
34 changed files with 3179 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/entrance.go
+4
View File
@@ -61,6 +61,7 @@ type APIGroup struct {
WafOwaspApi
WafHostPathRuleApi
WafAppApi
WafAIApi
}
var APIGroupAPP = new(APIGroup)
@@ -113,6 +114,9 @@ var (
wafAnalysisService = waf_service.WafAnalysisServiceApp
wafAIService = waf_service.WafAIServiceApp
wafAILabelService = waf_service.WafAILabelServiceApp
wafPrivateInfoService = waf_service.WafPrivateInfoServiceApp
wafPrivateGroupService = waf_service.WafPrivateGroupServiceApp
wafCacheRuleService = waf_service.WafCacheRuleServiceApp
api/waf_ai_api.go
+423
View File
@@ -0,0 +1,423 @@
package api
import (
"SamWaf/global"
"SamWaf/innerbean"
"SamWaf/model/common/response"
"SamWaf/model/request"
"SamWaf/utils"
"SamWaf/wafai"
"encoding/json"
"fmt"
"io"
"net/url"
"os"
"path/filepath"
"strings"
"sync/atomic"
"time"
"github.com/gin-gonic/gin"
)
type WafAIApi struct {
}
const (
aiModelDir = "ai_model"
aiModelFile = "current.swai"
aiExportDir = "ai_export"
maxAIModelUpload = 64 * 1024 * 1024 // 模型包上传上限 64MB
defaultExportMax = 200000 // 默认最多导出条数
)
// aiStatusResp AI 检测状态响应
type aiStatusResp struct {
GlobalEnable int64 `json:"global_enable"` // 全局总开关
Mode string `json:"mode"` // observe / block
ModelLoaded bool `json:"model_loaded"` // 当前是否已加载模型
FeatureVersion string `json:"feature_version"` // 引擎特征版本
Manifest *wafai.Manifest `json:"manifest"` // 当前模型元数据（未加载为 null）
}
// GetAIStatusApi 获取 AI 检测状态与当前模型信息
func (w *WafAIApi) GetAIStatusApi(c *gin.Context) {
resp := aiStatusResp{
GlobalEnable: global.GCONFIG_AI_ENABLE,
Mode: global.GCONFIG_AI_MODE,
FeatureVersion: wafai.FeatureVersion,
}
if global.GWAF_AI_DETECTOR != nil {
if m, ok := global.GWAF_AI_DETECTOR.CurrentManifest(); ok {
resp.ModelLoaded = true
resp.Manifest = &m
}
}
response.OkWithData(resp, c)
}
// UploadAIModelApi 上传 .swai 模型包：校验 -> 热加载 -> 持久化
func (w *WafAIApi) UploadAIModelApi(c *gin.Context) {
if global.GWAF_AI_DETECTOR == nil {
response.FailWithMessage("AI检测器未初始化", c)
return
}
file, err := c.FormFile("file")
if err != nil {
response.FailWithMessage("文件上传失败: "+err.Error(), c)
return
}
if strings.ToLower(filepath.Ext(file.Filename)) != ".swai" {
response.FailWithMessage("不支持的文件类型，仅支持 .swai 模型包", c)
return
}
if file.Size > maxAIModelUpload {
response.FailWithMessage("模型包过大（上限 64MB）", c)
return
}
src, err := file.Open()
if err != nil {
response.FailWithMessage("打开上传文件失败: "+err.Error(), c)
return
}
defer src.Close()
data, err := io.ReadAll(io.LimitReader(src, maxAIModelUpload+1))
if err != nil {
response.FailWithMessage("读取上传文件失败: "+err.Error(), c)
return
}
if int64(len(data)) > maxAIModelUpload {
response.FailWithMessage("模型包过大（上限 64MB）", c)
return
}
// 先校验+热加载到内存（含特征版本/sha256/zip 安全校验），通过后再落盘
manifest, err := global.GWAF_AI_DETECTOR.LoadFromBytes(data)
if err != nil {
response.FailWithMessage("模型校验失败: "+err.Error(), c)
return
}
// 持久化到 data/ai_model/current.swai（原子替换），供下次启动加载
dir := filepath.Join(utils.GetCurrentDir(), "data", aiModelDir)
if err = os.MkdirAll(dir, 0750); err != nil {
response.FailWithMessage("创建模型目录失败: "+err.Error(), c)
return
}
finalPath := filepath.Join(dir, aiModelFile)
tmpPath := finalPath + ".tmp"
if err = os.WriteFile(tmpPath, data, 0640); err != nil {
response.FailWithMessage("保存模型文件失败: "+err.Error(), c)
return
}
if err = os.Rename(tmpPath, finalPath); err != nil {
_ = os.Remove(tmpPath)
response.FailWithMessage("替换模型文件失败: "+err.Error(), c)
return
}
response.OkWithDetailed(manifest, "模型上传并加载成功", c)
}
// ReloadAIModelApi 从磁盘重新加载当前模型
func (w *WafAIApi) ReloadAIModelApi(c *gin.Context) {
if global.GWAF_AI_DETECTOR == nil {
response.FailWithMessage("AI检测器未初始化", c)
return
}
path := filepath.Join(utils.GetCurrentDir(), "data", aiModelDir, aiModelFile)
if _, err := os.Stat(path); err != nil {
response.FailWithMessage("没有可加载的模型文件，请先上传", c)
return
}
manifest, err := global.GWAF_AI_DETECTOR.LoadFromFile(path)
if err != nil {
response.FailWithMessage("模型加载失败: "+err.Error(), c)
return
}
response.OkWithDetailed(manifest, "模型重新加载成功", c)
}
// UnloadAIModelApi 卸载当前模型（不删除磁盘文件）
func (w *WafAIApi) UnloadAIModelApi(c *gin.Context) {
if global.GWAF_AI_DETECTOR != nil {
global.GWAF_AI_DETECTOR.Unload()
}
response.OkWithMessage("模型已卸载", c)
}
// GetAIDashboardApi AI检测看板：按类别汇总 + 分数分布 + observe/block 趋势
func (w *WafAIApi) GetAIDashboardApi(c *gin.Context) {
var req request.WafAIDashboardReq
_ = c.ShouldBindJSON(&req)
response.OkWithData(wafAIService.DashboardApi(req), c)
}
// MarkLabelApi 标记某条日志的训练标签修正（误报→正常 / 确认攻击 / 忽略）
func (w *WafAIApi) MarkLabelApi(c *gin.Context) {
var req request.WafAILabelMarkReq
if err := c.ShouldBindJSON(&req); err != nil {
response.FailWithMessage("参数错误: "+err.Error(), c)
return
}
if err := wafAILabelService.MarkApi(req); err != nil {
response.FailWithMessage("标记失败: "+err.Error(), c)
return
}
response.OkWithMessage("标记成功", c)
}
// UnmarkLabelApi 取消某条日志的标记
func (w *WafAIApi) UnmarkLabelApi(c *gin.Context) {
var req request.WafAILabelUnmarkReq
if err := c.ShouldBindJSON(&req); err != nil {
response.FailWithMessage("参数错误: "+err.Error(), c)
return
}
if err := wafAILabelService.UnmarkApi(req.ReqUuid); err != nil {
response.FailWithMessage("取消标记失败: "+err.Error(), c)
return
}
response.OkWithMessage("已取消标记", c)
}
// LabelByUuidsApi 按 req_uuid 批量查询标记状态（日志列表回显用）
func (w *WafAIApi) LabelByUuidsApi(c *gin.Context) {
var req request.WafAILabelByUuidsReq
_ = c.ShouldBindJSON(&req)
response.OkWithData(wafAILabelService.GetMapByUuidsApi(req.ReqUuids), c)
}
// LabelListApi 标注工作台列表：AI 命中分页 + 标记状态过滤/回显
func (w *WafAIApi) LabelListApi(c *gin.Context) {
var req request.WafAILabelListReq
_ = c.ShouldBindJSON(&req)
response.OkWithData(wafAILabelService.ListApi(req), c)
}
// BatchMarkLabelApi 批量标记训练标签（误报→正常 / 确认攻击 / 忽略）
func (w *WafAIApi) BatchMarkLabelApi(c *gin.Context) {
var req request.WafAILabelBatchMarkReq
if err := c.ShouldBindJSON(&req); err != nil {
response.FailWithMessage("参数错误: "+err.Error(), c)
return
}
n, err := wafAILabelService.BatchMarkApi(req)
if err != nil {
response.FailWithMessage(fmt.Sprintf("批量标记失败(已处理%d条): %s", n, err.Error()), c)
return
}
response.OkWithMessage(fmt.Sprintf("已标记 %d 条", n), c)
}
// BatchUnmarkLabelApi 批量取消标记
func (w *WafAIApi) BatchUnmarkLabelApi(c *gin.Context) {
var req request.WafAILabelBatchUnmarkReq
if err := c.ShouldBindJSON(&req); err != nil {
response.FailWithMessage("参数错误: "+err.Error(), c)
return
}
n, err := wafAILabelService.BatchUnmarkApi(req.ReqUuids)
if err != nil {
response.FailWithMessage("批量取消标记失败: "+err.Error(), c)
return
}
response.OkWithMessage(fmt.Sprintf("已取消 %d 条标记", n), c)
}
// trainSample 导出的训练样本（字段名与 SamWafAI samwafai/data/sample.py 对齐）
type trainSample struct {
Method string `json:"method"`
Path string `json:"path"`
Query string `json:"query"`
Body string `json:"body"`
UserAgent string `json:"user_agent"`
Label int `json:"label"`
AttackType string `json:"attack_type"`
Source string `json:"source"`
}
// ExportTrainDataApi 导出脱敏训练数据为 JSONL，落到 data/ai_export/ 供本地训练
// aiExportRunning 导出任务并发保护：同一时刻只允许一个导出在跑
var aiExportRunning int32
func (w *WafAIApi) ExportTrainDataApi(c *gin.Context) {
var req request.WafAIExportReq
_ = c.ShouldBindJSON(&req)
maxCount := req.MaxCount
if maxCount <= 0 || maxCount > defaultExportMax {
maxCount = defaultExportMax
}
if global.GWAF_LOCAL_LOG_DB == nil {
response.FailWithMessage("日志库未初始化", c)
return
}
// 导出是耗时操作（大表查询 + 逐行脱敏写盘），异步执行避免请求超时；
// 完成/失败通过 OpResultMessageInfo 推送到管理端通知中心。
if !atomic.CompareAndSwapInt32(&aiExportRunning, 0, 1) {
response.FailWithMessage("已有导出任务正在进行，请等待完成", c)
return
}
go func() {
defer atomic.StoreInt32(&aiExportRunning, 0)
outPath, nAttack, nNormal, nDrop, _, err := runAIExport(req, maxCount)
serverName := global.GWAF_CUSTOM_SERVER_NAME
if err != nil {
global.GQEQUE_MESSAGE_DB.Enqueue(innerbean.OpResultMessageInfo{
BaseMessageInfo: innerbean.BaseMessageInfo{OperaType: "AI训练数据导出", Server: serverName},
Msg: "AI训练数据导出失败：" + err.Error(),
Success: "false",
})
return
}
global.GQEQUE_MESSAGE_DB.Enqueue(innerbean.OpResultMessageInfo{
BaseMessageInfo: innerbean.BaseMessageInfo{OperaType: "AI训练数据导出", Server: serverName},
Msg: fmt.Sprintf("AI训练数据导出完毕：%s（写入%d条，攻击%d/正常%d，已忽略%d）",
outPath, nAttack+nNormal, nAttack, nNormal, nDrop),
Success: "true",
})
}()
response.OkWithMessage("导出任务已开始，完成后会在通知中心提示，文件生成在服务器 data/ai_export/ 目录", c)
}
// runAIExport 实际执行训练数据导出，返回文件路径与各类计数。
func runAIExport(req request.WafAIExportReq, maxCount int) (outPath string, nAttack, nNormal, nDrop, total int, err error) {
query := global.GWAF_LOCAL_LOG_DB.Model(&innerbean.WebLog{}).
Select("REQ_UUID", "METHOD", "URL", "RawQuery", "BODY", "POST_FORM", "USER_AGENT", "ACTION", "RULE", "LogOnlyMode")
if req.Days > 0 {
cutoff := time.Now().AddDate(0, 0, -req.Days).Format("2006-01-02 15:04:05")
query = query.Where("create_time >= ?", cutoff)
}
var rows []innerbean.WebLog
if err = query.Order("unix_add_time desc").Limit(maxCount).Find(&rows).Error; err != nil {
return "", 0, 0, 0, 0, fmt.Errorf("查询日志失败: %w", err)
}
total = len(rows)
dir := filepath.Join(utils.GetCurrentDir(), "data", aiExportDir)
if err = os.MkdirAll(dir, 0750); err != nil {
return "", 0, 0, 0, total, fmt.Errorf("创建导出目录失败: %w", err)
}
outPath = filepath.Join(dir, fmt.Sprintf("train_%s.jsonl", time.Now().Format("20060102_150405")))
fh, ferr := os.OpenFile(outPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0640)
if ferr != nil {
return "", 0, 0, 0, total, fmt.Errorf("创建导出文件失败: %w", ferr)
}
defer fh.Close()
enc := json.NewEncoder(fh)
// 一次性加载人工标记（含请求快照），人工修正优先于规则弱标签
marks := wafAILabelService.GetAllFull()
// 1) 过滤窗口内的日志：仅处理"未人工标记"的（已标记的统一在第2步从快照产出，
// 确保人工修正不会被时间/条数条件漏掉）
for i := range rows {
r := &rows[i]
if _, marked := marks[r.REQ_UUID]; marked {
continue
}
verdict, at := wafai.WeakLabel(r.ACTION, r.RULE, r.LogOnlyMode)
if verdict == wafai.VerdictDrop {
nDrop++
continue
}
label := 0
attackType := ""
if verdict == wafai.VerdictAttack {
// 仅信任高置信检测器(libinjection/OWASP)的攻击标签；
// 低置信(scan/rce/traversal 等)未经人工确认则丢弃，避免规则误报污染训练
if !wafai.IsHighConfidenceAttackType(at) {
nDrop++
continue
}
label = 1
attackType = at
}
path, rawQuery := splitURL(r.URL, r.RawQuery)
body := r.BODY
if body == "" {
body = r.POST_FORM
}
if label == 1 {
nAttack++
} else {
nNormal++
}
s := trainSample{
Method: strings.ToUpper(r.METHOD),
Path: path,
Query: desensitizeForExport(rawQuery),
Body: desensitizeForExport(body),
UserAgent: r.USER_AGENT,
Label: label,
AttackType: attackType,
Source: "samwaf_log",
}
if err = enc.Encode(&s); err != nil {
return outPath, nAttack, nNormal, nDrop, total, fmt.Errorf("写入导出文件失败: %w", err)
}
}
// 2) 全部人工标记从快照产出，不受时间/条数条件限制（连原始日志被清理也能产出）
for _, m := range marks {
if m.Mark == "ignore" {
nDrop++
continue
}
label := 0
if m.Mark == "attack" {
label = 1
}
path, rawQuery := splitURL(m.URL, m.RAW_QUERY)
if label == 1 {
nAttack++
} else {
nNormal++
}
s := trainSample{
Method: strings.ToUpper(m.METHOD),
Path: path,
Query: desensitizeForExport(rawQuery),
Body: desensitizeForExport(m.BODY),
UserAgent: m.USER_AGENT,
Label: label,
AttackType: m.AttackType,
Source: "manual",
}
if err = enc.Encode(&s); err != nil {
return outPath, nAttack, nNormal, nDrop, total, fmt.Errorf("写入导出文件失败: %w", err)
}
}
return outPath, nAttack, nNormal, nDrop, total, nil
}
// splitURL 将 URL 拆为 path 与 query；优先用已有的 RawQuery
func splitURL(rawURL, rawQuery string) (string, string) {
path := rawURL
if rawQuery == "" {
if u, err := url.Parse(rawURL); err == nil {
path = u.Path
rawQuery = u.RawQuery
}
} else if idx := strings.IndexByte(rawURL, '?'); idx >= 0 {
path = rawURL[:idx]
}
if path == "" {
path = "/"
}
return path, rawQuery
}
// desensitizeForExport 导出前再做一次脱敏（复用 godlp 引擎），避免敏感参数值随训练数据外泄
func desensitizeForExport(s string) string {
if s == "" || global.GWAF_DLP == nil {
return s
}
return utils.DeSenText(s)
}
cmd/samwaf/main.go
+13
View File
@@ -14,6 +14,7 @@ import (
"SamWaf/model/wafenginmodel"
"SamWaf/plugin"
"SamWaf/utils"
"SamWaf/wafai"
"SamWaf/wafappengine"
"SamWaf/wafconfig"
"SamWaf/wafdb"
@@ -263,6 +264,18 @@ func (m *wafSystenService) run() {
global.GWAF_DLP_CONFIG = ldpConfig
global.GWAF_REG_PUBLIC_KEY = publicKey
//初始化AI智能检测器，若存在已上传的模型包则加载（失败安全，不影响启动）
global.GWAF_AI_DETECTOR = wafai.NewDetector()
aiModelPath := filepath.Join(utils.GetCurrentDir(), "data", "ai_model", "current.swai")
if _, err := os.Stat(aiModelPath); err == nil {
if manifest, err := global.GWAF_AI_DETECTOR.LoadFromFile(aiModelPath); err != nil {
zlog.Warn("AI模型加载失败，AI检测将不可用: ", err.Error())
} else {
zlog.Info(fmt.Sprintf("AI模型加载成功: version=%s feature=%s type=%s",
manifest.ModelVersion, manifest.FeatureVersion, manifest.ModelType))
}
}
//owasp资源 释放
err = wafinit.CheckAndReleaseDataset(owaspAssets, utils.GetCurrentDir()+"/data/owasp", "owasp")
if err != nil {
global/config.go
+16 -12
View File
@@ -28,18 +28,22 @@ var (
GCONFIG_ENABLE_STRICT_IP_BINDING int64 = 1 // 是否启用严格IP绑定 1启用 0禁用
GCONFIG_ENABLE_REPLAY_PROTECT int64 = 1 // 防重放攻击开关 1启用 0禁用
GCONFIG_RECORD_ENABLE_OWASP int64 = 0 //启动OWASP数据检测
GCONFIG_OWASP_MODE string = "DetectionOnly" //OWASP 检测引擎工作模式: On(拦截) / DetectionOnly(观察/仅记录) / Off(关闭)
GCONFIG_OWASP_BLOCK_THRESHOLD int64 = 7 //OWASP 入站 anomaly score 阈值(官方默认 5,我们宽松到 7)
GCONFIG_RECORD_ENABLE_HTTP_80 int64 = 0 //启动80端口服务（为自动申请证书使用 HTTP文件验证类型，DNS验证不需要）
GCONFIG_RECORD_SSLOrder_EXPIRE_DAY int64 = 30 // 提前多少天进行自动申请
GCONFIG_RECORD_SSL_IP_CERT_IP string = "" // 获取IP证书时的IP地址
GCONFIG_RECORD_SSL_IP_EXPIRE_DAY int64 = 3 // IP证书提前多少天进行自动申请
GCONFIG_RECORD_SSLHTTP_CHECK int64 = 0 // ssl申请文件验证类型 是否校验原始路径HTTP响应代码 1 校验 0 不校验
GCONFIG_RECORD_SSLMinVerson string = "TLS 1.2" // ssl最低版本
GCONFIG_RECORD_SSLMaxVerson string = "TLS 1.3" // ssl最大版本
GCONFIG_RECORD_CONNECT_TIME_OUT int64 = 30 // 连接超时 默认30s
GCONFIG_RECORD_KEEPALIVE_TIME_OUT int64 = 30 // 保持活动超时 默认30s
GCONFIG_RECORD_ENABLE_OWASP int64 = 0 //启动OWASP数据检测
GCONFIG_OWASP_MODE string = "DetectionOnly" //OWASP 检测引擎工作模式: On(拦截) / DetectionOnly(观察/仅记录) / Off(关闭)
GCONFIG_AI_ENABLE int64 = 0 //AI智能检测总开关 1启用 0关闭（需先在AI模型管理上传模型包）
GCONFIG_AI_MODE string = "observe" //AI检测工作模式: observe(仅记录/观察) / block(达到拦截阈值则拦截)
GCONFIG_OWASP_BLOCK_THRESHOLD int64 = 7 //OWASP 入站 anomaly score 阈值(官方默认 5,我们宽松到 7)
GCONFIG_RECORD_ENABLE_HTTP_80 int64 = 0 //启动80端口服务（为自动申请证书使用 HTTP文件验证类型，DNS验证不需要）
GCONFIG_RECORD_SSLOrder_EXPIRE_DAY int64 = 30 // 提前多少天进行自动申请
GCONFIG_RECORD_SSL_IP_CERT_IP string = "" // 获取IP证书时的IP地址
GCONFIG_RECORD_SSL_IP_EXPIRE_DAY int64 = 3 // IP证书提前多少天进行自动申请
GCONFIG_RECORD_SSLHTTP_CHECK int64 = 0 // ssl申请文件验证类型 是否校验原始路径HTTP响应代码 1 校验 0 不校验
GCONFIG_RECORD_SSLMinVerson string = "TLS 1.2" // ssl最低版本
GCONFIG_RECORD_SSLMaxVerson string = "TLS 1.3" // ssl最大版本
GCONFIG_RECORD_CONNECT_TIME_OUT int64 = 30 // 连接超时 默认30s
GCONFIG_RECORD_KEEPALIVE_TIME_OUT int64 = 30 // 保持活动超时 默认30s
//GCONFIG_RECORD_PATCH_VERSION_CORE int64 = 20250106 // 核心数据库补丁日期
//GCONFIG_RECORD_PATCH_VERSION_LOG int64 = 20250106 // 日志数据库补丁日期
GCONFIG_RECORD_ALL_SRC_BYTE_INFO int64 = 0 //记录原始信息(默认不开启)
global/global.go
+2
View File
@@ -7,6 +7,7 @@ import (
"SamWaf/iplocation"
"SamWaf/model"
"SamWaf/model/spec"
"SamWaf/wafai"
"SamWaf/wafnotify"
"SamWaf/wafowasp"
"SamWaf/wafsnowflake"
@@ -99,6 +100,7 @@ var (
GWAF_APP_OP_PASSWORD string //应用操作密码（高危操作二次确认，自动生成存 config.yml）
GWAF_DLP dlpheader.EngineAPI // 脱敏引擎
GWAF_DLP_CONFIG string // 脱敏引擎配置数据
GWAF_AI_DETECTOR *wafai.Detector // AI智能检测器（持有当前模型，失败安全）
GWAF_OWASP *wafowasp.WafOWASP //owasp引擎（兼容保留：由 GWAF_OWASP_MANAGER.Current() 更新；请勿直接赋值新实例，否则热重载会失效）
GWAF_OWASP_MANAGER *wafowasp.OwaspManager //owasp 管理器（支持热重载）
go.mod
+1
View File
@@ -80,6 +80,7 @@ require (
github.com/cloudwego/base64x v0.1.6 // indirect
github.com/cyphar/filepath-securejoin v0.2.4 // indirect
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
github.com/dmitryikh/leaves v0.0.0-20230708180554-25d19a787328 // indirect
github.com/emirpasic/gods v1.18.1 // indirect
github.com/fatih/color v1.16.0 // indirect
github.com/fsnotify/fsnotify v1.9.0 // indirect
go.sum
+2
View File
@@ -111,6 +111,8 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/denisbrodbeck/machineid v1.0.1 h1:geKr9qtkB876mXguW2X6TU4ZynleN6ezuMSRhl4D7AQ=
github.com/denisbrodbeck/machineid v1.0.1/go.mod h1:dJUwb7PTidGDeYyUBmXZ2GphQBbjJCrnectwCyxcUSI=
github.com/dmitryikh/leaves v0.0.0-20230708180554-25d19a787328 h1:ht/zhLOAy9iiEKTKGkXvpw92Z7O6NK0bIVZVREy0kIE=
github.com/dmitryikh/leaves v0.0.0-20230708180554-25d19a787328/go.mod h1:wzMig9tMIJB8HsxXHppa9yRPo8BpO0eBM/Z4xnaohCQ=
github.com/dsnet/compress v0.0.1 h1:PlZu0n3Tuv04TzpfPbrnI0HW/YwodEXDS+oPKahKF0Q=
github.com/dsnet/compress v0.0.1/go.mod h1:Aw8dCMJ7RioblQeTqt88akK31OvO8Dhf5JflhBbQEHo=
github.com/dsnet/golib v0.0.0-20171103203638-1ea166775780/go.mod h1:Lj+Z9rebOhdfkVLjJ8T6VcRQv3SXugXy999NBtR9aFY=
innerbean/web_log.go
+50 -49
View File
@@ -3,55 +3,56 @@ package innerbean
import "strings"
type WebLog struct {
WafInnerDFlag string `gorm:"size:10" json:"waf_inner_dflag"` //日志队列处理方式
HOST string `gorm:"size:255" json:"host"`
URL string `gorm:"type:text" json:"url"`
RawQuery string `gorm:"type:text" json:"raw_query"` //原始URL查询
REFERER string `gorm:"type:text" json:"referer"`
USER_AGENT string `gorm:"size:500" json:"user_agent"`
METHOD string `gorm:"size:20" json:"method"`
HEADER string `gorm:"type:text" json:"header"`
SRC_IP string `gorm:"size:64" json:"src_ip"`
SRC_PORT string `gorm:"size:10" json:"src_port"`
COUNTRY string `gorm:"size:100" json:"country"`
PROVINCE string `gorm:"size:100" json:"province"`
CITY string `gorm:"size:100" json:"city"`
CREATE_TIME string `gorm:"size:32;index:idx_weblog_time" json:"create_time"`
CONTENT_LENGTH int64 `json:"content_length"`
RES_CONTENT_LENGTH int64 `json:"res_content_length"` //响应内容大小（字节）
COOKIES string `gorm:"type:text" json:"cookies"`
BODY string `gorm:"type:text" json:"body"`
REQ_UUID string `gorm:"size:64" json:"req_uuid"`
USER_CODE string `gorm:"size:64;index" json:"user_code"`
TenantId string `gorm:"size:64;index" json:"tenant_id"` //租户ID（主要键）
HOST_CODE string `gorm:"size:64" json:"host_code"` //主机ID （主要键）
Day int `json:"day"` //日 （主要键）
ACTION string `gorm:"size:100" json:"action"`
RULE string `gorm:"type:text" json:"rule"`
STATUS string `gorm:"size:50" json:"status"` //状态
STATUS_CODE int `json:"status_code"` //状态编码
RES_BODY string `gorm:"type:text" json:"res_body"` //返回信息
POST_FORM string `gorm:"type:text" json:"post_form"` //提交的表单数据
TASK_FLAG int `json:"task_flag" gorm:"default:-1;index"` //任务处理标记 -1 等待处理；1 可以进行处理；2 处理完毕
UNIX_ADD_TIME int64 `json:"unix_add_time" gorm:"index"` //添加日期unix
RISK_LEVEL int `json:"risk_level"` //危险等级 0:正常 1:轻微 2:有害 3:严重 4:特别严重
GUEST_IDENTIFICATION string `gorm:"column:guest_id_entification;size:191" json:"guest_identification"` //访客身份识别
IsBot int `json:"is_bot"` //是否是机器人 0 不是机器人 1 机器人
TimeSpent int64 `json:"time_spent"` //用时
NetSrcIp string `gorm:"size:64" json:"net_src_ip"` //获取的原始IP
SrcByteBody []byte `json:"src_byte_body"` //原始body信息
SrcByteResBody []byte `json:"src_byte_res_body"` //返回body bytes信息
WebLogVersion int `json:"web_log_version"` //日志版本信息早期的是空和0，后期实时增加
Scheme string `gorm:"size:20" json:"scheme"` //HTTP 协议
SrcURL []byte `json:"src_url"` //原始url信息
PreCheckCost int64 `json:"pre_check_cost"` // 前置检查耗时(ms)
ForwardCost int64 `json:"forward_cost"` // 转发耗时(ms)
BackendCheckCost int64 `json:"backend_check_cost"` // 后端处理耗时(ms)
ResHeader string `gorm:"type:text" json:"res_header"` // 返回header情况
BodyHash string `gorm:"size:100" json:"body_hash"` // body hash值
LogOnlyMode int `json:"log_only_mode"` //是否只记录日志 1 是 0 不是
IsBalance int `json:"is_balance"` //是否是负载均衡 1 是 0 不是
BalanceInfo string `gorm:"size:255" json:"balance_info"` //负载均衡IP端口信息
WafInnerDFlag string `gorm:"size:10" json:"waf_inner_dflag"` //日志队列处理方式
HOST string `gorm:"size:255" json:"host"`
URL string `gorm:"type:text" json:"url"`
RawQuery string `gorm:"type:text" json:"raw_query"` //原始URL查询
REFERER string `gorm:"type:text" json:"referer"`
USER_AGENT string `gorm:"size:500" json:"user_agent"`
METHOD string `gorm:"size:20" json:"method"`
HEADER string `gorm:"type:text" json:"header"`
SRC_IP string `gorm:"size:64" json:"src_ip"`
SRC_PORT string `gorm:"size:10" json:"src_port"`
COUNTRY string `gorm:"size:100" json:"country"`
PROVINCE string `gorm:"size:100" json:"province"`
CITY string `gorm:"size:100" json:"city"`
CREATE_TIME string `gorm:"size:32;index:idx_weblog_time" json:"create_time"`
CONTENT_LENGTH int64 `json:"content_length"`
RES_CONTENT_LENGTH int64 `json:"res_content_length"` //响应内容大小（字节）
COOKIES string `gorm:"type:text" json:"cookies"`
BODY string `gorm:"type:text" json:"body"`
REQ_UUID string `gorm:"size:64" json:"req_uuid"`
USER_CODE string `gorm:"size:64;index" json:"user_code"`
TenantId string `gorm:"size:64;index" json:"tenant_id"` //租户ID（主要键）
HOST_CODE string `gorm:"size:64" json:"host_code"` //主机ID （主要键）
Day int `json:"day"` //日 （主要键）
ACTION string `gorm:"size:100" json:"action"`
RULE string `gorm:"type:text" json:"rule"`
STATUS string `gorm:"size:50" json:"status"` //状态
STATUS_CODE int `json:"status_code"` //状态编码
RES_BODY string `gorm:"type:text" json:"res_body"` //返回信息
POST_FORM string `gorm:"type:text" json:"post_form"` //提交的表单数据
TASK_FLAG int `json:"task_flag" gorm:"default:-1;index"` //任务处理标记 -1 等待处理；1 可以进行处理；2 处理完毕
UNIX_ADD_TIME int64 `json:"unix_add_time" gorm:"index"` //添加日期unix
RISK_LEVEL int `json:"risk_level"` //危险等级 0:正常 1:轻微 2:有害 3:严重 4:特别严重
GUEST_IDENTIFICATION string `gorm:"column:guest_id_entification;size:191" json:"guest_identification"` //访客身份识别
IsBot int `json:"is_bot"` //是否是机器人 0 不是机器人 1 机器人
TimeSpent int64 `json:"time_spent"` //用时
NetSrcIp string `gorm:"size:64" json:"net_src_ip"` //获取的原始IP
SrcByteBody []byte `json:"src_byte_body"` //原始body信息
SrcByteResBody []byte `json:"src_byte_res_body"` //返回body bytes信息
WebLogVersion int `json:"web_log_version"` //日志版本信息早期的是空和0，后期实时增加
Scheme string `gorm:"size:20" json:"scheme"` //HTTP 协议
SrcURL []byte `json:"src_url"` //原始url信息
PreCheckCost int64 `json:"pre_check_cost"` // 前置检查耗时(ms)
ForwardCost int64 `json:"forward_cost"` // 转发耗时(ms)
BackendCheckCost int64 `json:"backend_check_cost"` // 后端处理耗时(ms)
ResHeader string `gorm:"type:text" json:"res_header"` // 返回header情况
BodyHash string `gorm:"size:100" json:"body_hash"` // body hash值
LogOnlyMode int `json:"log_only_mode"` //是否只记录日志 1 是 0 不是
IsBalance int `json:"is_balance"` //是否是负载均衡 1 是 0 不是
BalanceInfo string `gorm:"size:255" json:"balance_info"` //负载均衡IP端口信息
AI_SCORE float64 `json:"ai_score"` //AI检测得分[0,1]，0表示未经AI检测或未命中；命中(观察/拦截)时记录实际分数
}
// GetHeaderValue 从HEADER字段中提取指定header的值
model/hosts.go
+2
View File
@@ -61,6 +61,7 @@ type HostsDefense struct {
DEFENSE_SENSITIVE int `json:"sensitive"` //敏感词检测
DEFENSE_DIR_TRAVERSAL int `json:"traversal"` //目录穿越检测
DEFENSE_OWASP_SET int `json:"owaspset"` //OWASP集检测
DEFENSE_AI int `json:"ai"` //AI智能检测（默认关闭，需先上传模型包并开启全局AI开关）
}
// HealthyConfig 健康度检测
@@ -307,6 +308,7 @@ func ParseHostsDefense(defenseJSON string) HostsDefense {
defense.DEFENSE_SENSITIVE = 1
defense.DEFENSE_DIR_TRAVERSAL = 1
defense.DEFENSE_OWASP_SET = 0
defense.DEFENSE_AI = 0
// 如果JSON不为空，则解析覆盖默认值
if defenseJSON != "" {
model/request/waf_ai_req.go
+58
View File
@@ -0,0 +1,58 @@
package request
// WafAIExportReq 训练数据导出入参
type WafAIExportReq struct {
Days int `json:"days"` // 导出最近多少天的日志（<=0 表示不限）
MaxCount int `json:"max_count"` // 最多导出多少条（<=0 表示用默认上限）
}
// WafAIDashboardReq AI检测看板入参
type WafAIDashboardReq struct {
StartDay int `json:"start_day"` // 起始日 YYYYMMDD（<=0 表示不限）
EndDay int `json:"end_day"` // 结束日 YYYYMMDD（<=0 表示不限）
HostCode string `json:"host_code"` // 站点编码，空表示全部站点
}
// WafAILabelMarkReq 训练标签人工修正入参
type WafAILabelMarkReq struct {
ReqUuid string `json:"req_uuid" binding:"required"` // 关联请求日志
HostCode string `json:"host_code"`
Mark string `json:"mark" binding:"required"` // normal / attack / ignore
AttackType string `json:"attack_type"` // mark=attack 时的人工分类，空=自动判定
Rule string `json:"rule"` // 原始触发规则（快照）
SrcIp string `json:"src_ip"`
Url string `json:"url"`
}
// WafAILabelUnmarkReq 取消标记入参
type WafAILabelUnmarkReq struct {
ReqUuid string `json:"req_uuid" binding:"required"`
}
// WafAILabelByUuidsReq 按 req_uuid 批量查询标记状态（用于日志列表回显）
type WafAILabelByUuidsReq struct {
ReqUuids []string `json:"req_uuids"`
}
// WafAILabelListReq 标注工作台列表入参：在 AI 命中(ai_score>0)子集上分页 + 过滤
type WafAILabelListReq struct {
HostCode string `json:"host_code"` // 站点编码，空=全部
StartDay int `json:"start_day"` // 起始日 YYYYMMDD（<=0 不限）
EndDay int `json:"end_day"` // 结束日 YYYYMMDD（<=0 不限）
MarkStatus string `json:"mark_status"` // ""=全部 / unmarked / marked / normal / attack / ignore
MinScore float64 `json:"min_score"` // 最小分数（<=0 不限）
PageIndex int `json:"page_index"` // 页码，从 1 开始
PageSize int `json:"page_size"` // 每页条数
}
// WafAILabelBatchMarkReq 批量标记入参
type WafAILabelBatchMarkReq struct {
ReqUuids []string `json:"req_uuids" binding:"required"` // 待标记的请求列表
Mark string `json:"mark" binding:"required"` // normal / attack / ignore
AttackType string `json:"attack_type"` // mark=attack 时的人工分类，空=每条按原始日志自动判定
}
// WafAILabelBatchUnmarkReq 批量取消标记入参
type WafAILabelBatchUnmarkReq struct {
ReqUuids []string `json:"req_uuids" binding:"required"`
}
model/response/waf_ai_resp.go
+48
View File
@@ -0,0 +1,48 @@
package response
// WafAINameValue 通用名称-数值对（类别汇总/分数分布用）
type WafAINameValue struct {
Name string `json:"name"`
Value int64 `json:"value"`
}
// WafAITrendPoint AI命中按天趋势点
type WafAITrendPoint struct {
Day int `json:"day"` // YYYYMMDD
Observe int64 `json:"observe"` // 观察命中数（log_only_mode=1）
Block int64 `json:"block"` // 拦截命中数（log_only_mode=0）
}
// WafAIDashboard AI检测看板聚合结果
type WafAIDashboard struct {
Total int64 `json:"total"` // AI命中总数
ObserveCnt int64 `json:"observe_cnt"` // 观察命中数
BlockCnt int64 `json:"block_cnt"` // 拦截命中数
Categories []WafAINameValue `json:"categories"` // 按类别汇总
ScoreHist []WafAINameValue `json:"score_hist"` // 分数分布直方图（10桶）
Trend []WafAITrendPoint `json:"trend"` // 按天 observe/block 趋势
}
// WafAILabelItem 标注工作台列表项（AI命中 + 当前标记状态 + 请求详情）
type WafAILabelItem struct {
ReqUuid string `json:"req_uuid"`
CreateTime string `json:"create_time"`
HostCode string `json:"host_code"`
SrcIp string `json:"src_ip"`
Method string `json:"method"`
Url string `json:"url"`
RawQuery string `json:"raw_query"`
Body string `json:"body"`
UserAgent string `json:"user_agent"`
AiScore float64 `json:"ai_score"`
Rule string `json:"rule"` // 命中规则文本（AI检测:<类别>）
LogOnlyMode int `json:"log_only_mode"` // 1 观察 0 拦截
Mark string `json:"mark"` // 当前人工标记 normal/attack/ignore，空=未标记
AttackType string `json:"attack_type"` // 人工分类
}
// WafAILabelList 标注工作台分页结果
type WafAILabelList struct {
Total int64 `json:"total"`
Rows []WafAILabelItem `json:"rows"`
}
model/waf_log_label_mark.go
+25
View File
@@ -0,0 +1,25 @@
package model
import "SamWaf/model/baseorm"
/*
*
训练标签人工修正
针对单条访问日志（按 req_uuid）修正其训练标签，用于纠正规则误报导致的错误弱标签。
导出训练数据时优先采用此处的人工修正结果。
*/
type WafLogLabelMark struct {
baseorm.BaseOrm
REQ_UUID string `gorm:"size:64;index" json:"req_uuid"` // 关联的请求日志
HOST_CODE string `gorm:"size:64" json:"host_code"` // 主机码
Mark string `gorm:"size:20" json:"mark"` // normal=实际正常 / attack=确认攻击 / ignore=不参与训练
AttackType string `gorm:"size:32" json:"attack_type"` // 当 mark=attack 时的人工分类(sqli/xss/rce/traversal/inject/scan/other)，空=用自动判定
RULE string `gorm:"type:text" json:"rule"` // 原始触发规则（快照，便于审阅）
SRC_IP string `gorm:"size:64" json:"src_ip"` // 来源IP（快照）
URL string `gorm:"type:text" json:"url"` // 访问URL（快照）
// 训练字段快照：标记时即固化，导出不受时间/条数条件影响，日志被清理也不丢
METHOD string `gorm:"size:20" json:"method"` // 请求方法（快照）
RAW_QUERY string `gorm:"type:text" json:"raw_query"` // 查询串（快照）
BODY string `gorm:"type:text" json:"body"` // 请求体（快照，body 或 post_form）
USER_AGENT string `gorm:"size:500" json:"user_agent"` // UA（快照）
}
router/entrance.go
+1
View File
@@ -59,6 +59,7 @@ type ApiGroup struct {
WafOwaspRouter
WafHostPathRuleRouter
WafAppRouter
WafAIRouter
}
type PublicApiGroup struct {
LoginRouter
router/waf_ai_router.go
+32
View File
@@ -0,0 +1,32 @@
package router
import (
"SamWaf/api"
"github.com/gin-gonic/gin"
)
type WafAIRouter struct {
}
// InitWafAIRouter AI智能检测：模型管理与训练数据导出。
// 涉及模型上传（会被引擎加载）与数据导出（数据出库），属安全敏感接口，
// 应挂在仅 Token 登录可访问的路由组上。
func (receiver *WafAIRouter) InitWafAIRouter(group *gin.RouterGroup) {
apiInstance := api.APIGroupAPP.WafAIApi
router := group.Group("/api/v1/ai")
{
router.GET("/status", apiInstance.GetAIStatusApi)
router.POST("/dashboard", apiInstance.GetAIDashboardApi)
router.POST("/model/upload", apiInstance.UploadAIModelApi)
router.POST("/model/reload", apiInstance.ReloadAIModelApi)
router.POST("/model/unload", apiInstance.UnloadAIModelApi)
router.POST("/export", apiInstance.ExportTrainDataApi)
router.POST("/label/mark", apiInstance.MarkLabelApi)
router.POST("/label/unmark", apiInstance.UnmarkLabelApi)
router.POST("/label/by_uuids", apiInstance.LabelByUuidsApi)
router.POST("/label/list", apiInstance.LabelListApi)
router.POST("/label/batch_mark", apiInstance.BatchMarkLabelApi)
router.POST("/label/batch_unmark", apiInstance.BatchUnmarkLabelApi)
}
}
service/waf_service/waf_ai_label_service.go
+316
View File
@@ -0,0 +1,316 @@
package waf_service
import (
"SamWaf/common/uuid"
"SamWaf/customtype"
"SamWaf/global"
"SamWaf/innerbean"
"SamWaf/model"
"SamWaf/model/baseorm"
"SamWaf/model/request"
response2 "SamWaf/model/response"
"SamWaf/wafai"
"errors"
"strings"
"time"
"gorm.io/gorm"
)
type WafAILabelService struct{}
var WafAILabelServiceApp = new(WafAILabelService)
// 合法的标记取值
var validMarks = map[string]bool{"normal": true, "attack": true, "ignore": true}
// LabelMarkInfo 标记回显/导出用的精简信息
type LabelMarkInfo struct {
Mark string `json:"mark"`
AttackType string `json:"attack_type"`
}
// MarkApi 新增/更新某条日志的训练标签修正（按 req_uuid 幂等 upsert）。
// 标记时即从日志库读取该请求的训练字段做快照，导出不再受时间/条数条件限制，
// 即使原始日志后续被清理也能产出该样本。
func (receiver *WafAILabelService) MarkApi(req request.WafAILabelMarkReq) error {
if !validMarks[req.Mark] {
return errors.New("非法的标记类型")
}
// 读取原始日志做快照（标记时日志通常仍存在）
var wl innerbean.WebLog
global.GWAF_LOCAL_LOG_DB.
Select("METHOD", "URL", "RawQuery", "BODY", "POST_FORM", "USER_AGENT", "ACTION", "RULE", "SRC_IP", "HOST_CODE", "LogOnlyMode").
Where("REQ_UUID = ?", req.ReqUuid).Limit(1).Find(&wl)
body := wl.BODY
if body == "" {
body = wl.POST_FORM
}
hostCode := wl.HOST_CODE
if hostCode == "" {
hostCode = req.HostCode
}
rule := wl.RULE
if rule == "" {
rule = req.Rule
}
srcIp := wl.SRC_IP
if srcIp == "" {
srcIp = req.SrcIp
}
url := wl.URL
if url == "" {
url = req.Url
}
// 解析攻击分类：优先人工指定；为空且标记为攻击时按原始日志自动判定
attackType := req.AttackType
if req.Mark == "attack" && attackType == "" {
if _, at := wafai.WeakLabel(wl.ACTION, wl.RULE, wl.LogOnlyMode); at != "" {
attackType = at
} else {
attackType = "other"
}
}
fields := map[string]interface{}{
"mark": req.Mark,
"attack_type": attackType,
"rule": rule,
"host_code": hostCode,
"src_ip": srcIp,
"url": url,
"method": strings.ToUpper(wl.METHOD),
"raw_query": wl.RawQuery,
"body": body,
"user_agent": wl.USER_AGENT,
"update_time": customtype.JsonTime(time.Now()),
}
var existing model.WafLogLabelMark
global.GWAF_LOCAL_DB.Where("req_uuid = ? and tenant_id = ? and user_code = ?",
req.ReqUuid, global.GWAF_TENANT_ID, global.GWAF_USER_CODE).First(&existing)
if existing.Id != "" {
return global.GWAF_LOCAL_DB.Model(&model.WafLogLabelMark{}).Where("id = ?", existing.Id).
Updates(fields).Error
}
bean := &model.WafLogLabelMark{
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()),
},
REQ_UUID: req.ReqUuid,
HOST_CODE: hostCode,
Mark: req.Mark,
AttackType: attackType,
RULE: rule,
SRC_IP: srcIp,
URL: url,
METHOD: strings.ToUpper(wl.METHOD),
RAW_QUERY: wl.RawQuery,
BODY: body,
USER_AGENT: wl.USER_AGENT,
}
return global.GWAF_LOCAL_DB.Create(bean).Error
}
// UnmarkApi 取消某条日志的标记
func (receiver *WafAILabelService) UnmarkApi(reqUuid string) error {
return global.GWAF_LOCAL_DB.Where("req_uuid = ? and tenant_id = ? and user_code = ?",
reqUuid, global.GWAF_TENANT_ID, global.GWAF_USER_CODE).Delete(&model.WafLogLabelMark{}).Error
}
// GetMapByUuidsApi 按 req_uuid 批量返回标记（用于日志列表回显）
func (receiver *WafAILabelService) GetMapByUuidsApi(uuids []string) map[string]LabelMarkInfo {
result := map[string]LabelMarkInfo{}
if len(uuids) == 0 {
return result
}
var rows []model.WafLogLabelMark
global.GWAF_LOCAL_DB.Where("tenant_id = ? and user_code = ? and req_uuid in ?",
global.GWAF_TENANT_ID, global.GWAF_USER_CODE, uuids).Find(&rows)
for _, r := range rows {
result[r.REQ_UUID] = LabelMarkInfo{Mark: r.Mark, AttackType: r.AttackType}
}
return result
}
// GetAllFull 返回当前租户/用户的全部标记（含请求快照），导出时一次性加载，按 req_uuid 索引。
func (receiver *WafAILabelService) GetAllFull() map[string]model.WafLogLabelMark {
result := map[string]model.WafLogLabelMark{}
var rows []model.WafLogLabelMark
global.GWAF_LOCAL_DB.Where("tenant_id = ? and user_code = ?",
global.GWAF_TENANT_ID, global.GWAF_USER_CODE).Find(&rows)
for _, r := range rows {
result[r.REQ_UUID] = r
}
return result
}
// GetMarkStatusMap 轻量返回当前租户/用户全部标记的状态（mark+attack_type），不含请求快照。
// 供标注工作台列表回显与按标记状态过滤使用。
func (receiver *WafAILabelService) GetMarkStatusMap() map[string]LabelMarkInfo {
result := map[string]LabelMarkInfo{}
var rows []model.WafLogLabelMark
global.GWAF_LOCAL_DB.Select("req_uuid", "mark", "attack_type").
Where("tenant_id = ? and user_code = ?", global.GWAF_TENANT_ID, global.GWAF_USER_CODE).Find(&rows)
for _, r := range rows {
result[r.REQ_UUID] = LabelMarkInfo{Mark: r.Mark, AttackType: r.AttackType}
}
return result
}
// ListApi 标注工作台列表：在 AI 命中(ai_score>0)子集上分页查询，并合并人工标记状态。
// 标记数据在 core 库、日志在 log 库（跨库无法 JOIN），故先取本租户标记集合，
// 在内存中做"按标记状态过滤"与逐行回显。
func (receiver *WafAILabelService) ListApi(req request.WafAILabelListReq) response2.WafAILabelList {
res := response2.WafAILabelList{Rows: []response2.WafAILabelItem{}}
if global.GWAF_LOCAL_LOG_DB == nil {
return res
}
pageIndex := req.PageIndex
if pageIndex <= 0 {
pageIndex = 1
}
pageSize := req.PageSize
if pageSize <= 0 || pageSize > 200 {
pageSize = 20
}
markMap := receiver.GetMarkStatusMap()
// 按标记状态预先算出 IN / NOT IN 的 uuid 集合
var inUuids []string // 命中即纳入
useNotIn := false // true 时改用 NOT IN
emptyResult := false // 过滤后必然为空，直接返回
switch req.MarkStatus {
case "unmarked":
for u := range markMap {
inUuids = append(inUuids, u)
}
useNotIn = true // 排除全部已标记
case "marked":
for u := range markMap {
inUuids = append(inUuids, u)
}
if len(inUuids) == 0 {
emptyResult = true
}
case "normal", "attack", "ignore":
for u, info := range markMap {
if info.Mark == req.MarkStatus {
inUuids = append(inUuids, u)
}
}
if len(inUuids) == 0 {
emptyResult = true
}
}
if emptyResult {
return res
}
buildQ := func() *gorm.DB {
q := global.GWAF_LOCAL_LOG_DB.Model(&innerbean.WebLog{}).Where("ai_score > 0")
if req.StartDay > 0 && req.EndDay > 0 {
q = q.Where("day between ? and ?", req.StartDay, req.EndDay)
}
if req.HostCode != "" {
q = q.Where("host_code = ?", req.HostCode)
}
if req.MinScore > 0 {
q = q.Where("ai_score >= ?", req.MinScore)
}
if len(inUuids) > 0 {
if useNotIn {
q = q.Where("req_uuid not in ?", inUuids)
} else {
q = q.Where("req_uuid in ?", inUuids)
}
}
return q
}
buildQ().Count(&res.Total)
var rows []innerbean.WebLog
buildQ().Select("REQ_UUID", "CREATE_TIME", "HOST_CODE", "SRC_IP", "METHOD", "URL",
"RawQuery", "BODY", "POST_FORM", "USER_AGENT", "AI_SCORE", "RULE", "LogOnlyMode").
Order("ai_score desc").Order("unix_add_time desc").
Offset((pageIndex - 1) * pageSize).Limit(pageSize).Find(&rows)
for i := range rows {
r := &rows[i]
body := r.BODY
if body == "" {
body = r.POST_FORM
}
item := response2.WafAILabelItem{
ReqUuid: r.REQ_UUID,
CreateTime: r.CREATE_TIME,
HostCode: r.HOST_CODE,
SrcIp: r.SRC_IP,
Method: strings.ToUpper(r.METHOD),
Url: r.URL,
RawQuery: r.RawQuery,
Body: body,
UserAgent: r.USER_AGENT,
AiScore: r.AI_SCORE,
Rule: r.RULE,
LogOnlyMode: r.LogOnlyMode,
}
if info, ok := markMap[r.REQ_UUID]; ok {
item.Mark = info.Mark
item.AttackType = info.AttackType
}
res.Rows = append(res.Rows, item)
}
return res
}
// BatchMarkApi 批量标记：逐条复用 MarkApi（每条独立读取请求快照、按需自动判定分类）。
// 返回成功标记条数。
func (receiver *WafAILabelService) BatchMarkApi(req request.WafAILabelBatchMarkReq) (int, error) {
if !validMarks[req.Mark] {
return 0, errors.New("非法的标记类型")
}
n := 0
for _, u := range req.ReqUuids {
if strings.TrimSpace(u) == "" {
continue
}
if err := receiver.MarkApi(request.WafAILabelMarkReq{
ReqUuid: u,
Mark: req.Mark,
AttackType: req.AttackType,
}); err != nil {
return n, err
}
n++
}
return n, nil
}
// BatchUnmarkApi 批量取消标记，返回删除条数。
func (receiver *WafAILabelService) BatchUnmarkApi(uuids []string) (int, error) {
var clean []string
for _, u := range uuids {
if strings.TrimSpace(u) != "" {
clean = append(clean, u)
}
}
if len(clean) == 0 {
return 0, nil
}
res := global.GWAF_LOCAL_DB.Where("tenant_id = ? and user_code = ? and req_uuid in ?",
global.GWAF_TENANT_ID, global.GWAF_USER_CODE, clean).Delete(&model.WafLogLabelMark{})
return int(res.RowsAffected), res.Error
}
service/waf_service/waf_ai_service.go
+92
View File
@@ -0,0 +1,92 @@
package waf_service
import (
"SamWaf/global"
"SamWaf/innerbean"
"SamWaf/model/request"
response2 "SamWaf/model/response"
"fmt"
"gorm.io/gorm"
)
type WafAIService struct{}
var WafAIServiceApp = new(WafAIService)
// DashboardApi 聚合 AI 检测看板数据：按类别汇总、分数分布、observe/block 趋势。
// 数据源为 web_logs 中 ai_score>0 的命中子集（相对较小），按 day 范围/站点过滤。
func (receiver *WafAIService) DashboardApi(req request.WafAIDashboardReq) response2.WafAIDashboard {
var res response2.WafAIDashboard
res.Categories = []response2.WafAINameValue{}
res.Trend = []response2.WafAITrendPoint{}
if global.GWAF_LOCAL_LOG_DB == nil {
res.ScoreHist = buildEmptyScoreHist()
return res
}
// 每次查询都用全新的 where 链，避免 GORM 条件被复用污染
base := func() *gorm.DB {
q := global.GWAF_LOCAL_LOG_DB.Model(&innerbean.WebLog{}).Where("ai_score > 0")
if req.StartDay > 0 && req.EndDay > 0 {
q = q.Where("day between ? and ?", req.StartDay, req.EndDay)
}
if req.HostCode != "" {
q = q.Where("host_code = ?", req.HostCode)
}
return q
}
// 1) 按类别（rule）汇总
base().Select("rule as name, count(*) as value").
Group("rule").Order("value desc").Scan(&res.Categories)
// 2) 分数分布直方图（10 桶：0.0-0.1 ... 0.9-1.0；score==1.0 归入最后一桶）
type bucketRow struct {
Bucket int
Cnt int64
}
var brows []bucketRow
base().Select("cast(ai_score*10 as int) as bucket, count(*) as cnt").
Group("bucket").Scan(&brows)
counts := make([]int64, 10)
for _, b := range brows {
idx := b.Bucket
if idx >= 10 {
idx = 9 // score==1.0 归入 0.9-1.0
}
if idx < 0 {
idx = 0
}
counts[idx] += b.Cnt
}
res.ScoreHist = make([]response2.WafAINameValue, 10)
for i := 0; i < 10; i++ {
res.ScoreHist[i] = response2.WafAINameValue{
Name: fmt.Sprintf("%.1f-%.1f", float64(i)/10, float64(i+1)/10),
Value: counts[i],
}
}
// 3) 按天 observe/block 趋势
base().Select("day, sum(case when log_only_mode=1 then 1 else 0 end) as observe, sum(case when log_only_mode=0 then 1 else 0 end) as block").
Group("day").Order("day asc").Scan(&res.Trend)
// 4) 汇总总数
for _, t := range res.Trend {
res.ObserveCnt += t.Observe
res.BlockCnt += t.Block
}
res.Total = res.ObserveCnt + res.BlockCnt
return res
}
func buildEmptyScoreHist() []response2.WafAINameValue {
hist := make([]response2.WafAINameValue, 10)
for i := 0; i < 10; i++ {
hist[i] = response2.WafAINameValue{Name: fmt.Sprintf("%.1f-%.1f", float64(i)/10, float64(i+1)/10)}
}
return hist
}
wafai/engine.go
+14
View File
@@ -0,0 +1,14 @@
package wafai
// InferenceEngine 推理引擎抽象。
//
// 一期实现为纯 Go GBDT（engine_gbdt.go，dmitryikh/leaves 加载 LightGBM）。
// 二期可选 ONNX 深度模型增强档（engine_onnx.go，build tag 隔离）。
type InferenceEngine interface {
// Score 输入特征向量，返回 [0,1] 的攻击概率。
Score(features []float64) float64
// FeatureVersion 模型训练时使用的特征版本，用于与 Go 侧特征版本校验。
FeatureVersion() string
// Type 引擎类型标识："gbdt" / "onnx"。
Type() string
}
wafai/engine_gbdt.go
+38
View File
@@ -0,0 +1,38 @@
package wafai
import (
"bufio"
"bytes"
"fmt"
"github.com/dmitryikh/leaves"
)
// gbdtEngine 基于 dmitryikh/leaves 的纯 Go LightGBM 推理引擎（一期）。
type gbdtEngine struct {
model *leaves.Ensemble
featureVersion string
}
// newGBDTEngine 从 LightGBM 文本模型字节构建引擎。
// loadTransformation=true 让 leaves 加载 sigmoid 变换，PredictSingle 直接输出概率。
func newGBDTEngine(modelBytes []byte, featureVersion string) (*gbdtEngine, error) {
br := bufio.NewReader(bytes.NewReader(modelBytes))
model, err := leaves.LGEnsembleFromReader(br, true)
if err != nil {
return nil, fmt.Errorf("加载 LightGBM 模型失败: %w", err)
}
if nf := model.NFeatures(); nf != FeatureCount {
return nil, fmt.Errorf("模型特征维度=%d 与引擎=%d 不一致", nf, FeatureCount)
}
return &gbdtEngine{model: model, featureVersion: featureVersion}, nil
}
func (e *gbdtEngine) Score(features []float64) float64 {
// nIterations=0 表示使用全部树
return e.model.PredictSingle(features, 0)
}
func (e *gbdtEngine) FeatureVersion() string { return e.featureVersion }
func (e *gbdtEngine) Type() string { return "gbdt" }
wafai/engine_onnx_stub.go
+11
View File
@@ -0,0 +1,11 @@
//go:build !samwaf_onnx
package wafai
import "errors"
// newONNXEngine 默认构建不含 ONNX 运行时（保持主分支零新增 CGO/外部依赖）。
// 需要 ONNX 深度模型增强档时，用 `-tags samwaf_onnx` 编译并提供 engine_onnx.go 实现。
func newONNXEngine(modelBytes []byte, featureVersion string) (InferenceEngine, error) {
return nil, errors.New("当前构建未启用 ONNX 引擎（需使用 -tags samwaf_onnx 编译）")
}
wafai/features.go
+394
View File
@@ -0,0 +1,394 @@
package wafai
import (
"bytes"
"math"
)
// 特征规范 v1
//
// !!! 重要 !!!
// 本文件与 SamWafAI 仓库 samwafai/features/lexical.py 是同一份规范
// （SamWafTechDoc 特征规范 v1）的双语言实现，二者必须逐字节对齐。
// 任何改动（常量/词表/顺序/算法）必须同步修改 Python 侧并升级 FeatureVersion，
// 且重新生成 wafai/testdata/feature_golden.json（uv run samwafai gen-golden）。
const (
FeatureVersion = "v1"
maxComponentBytes = 4096 // 每个分量截断长度
decodeMax = 3 // 百分号解码最大迭代次数
ngramN = 3 // 字节 n-gram 长度
ngramBuckets = 64 // n-gram 哈希桶数量
scalarFeatureCount = 22
FeatureCount = scalarFeatureCount + ngramBuckets // 86
)
// methodIdx HTTP 方法编码（大写比较；未知方法 -> 7）
var methodIdx = map[string]float64{
"GET": 0, "POST": 1, "PUT": 2, "DELETE": 3,
"HEAD": 4, "OPTIONS": 5, "PATCH": 6,
}
const methodIdxOther = 7
// 关键词表（小写；与 Python 侧严格一致）
var (
kwSQL = [][]byte{
[]byte("select"), []byte("union"), []byte("insert into"), []byte("update "),
[]byte("delete from"), []byte("drop table"), []byte("information_schema"),
[]byte("sleep("), []byte("benchmark("), []byte("load_file"), []byte("group by"),
[]byte("order by"), []byte("or 1=1"), []byte("' or '"), []byte("\" or \""),
[]byte("concat("), []byte("char("), []byte("0x"), []byte("xp_"), []byte("exec("),
[]byte("waitfor delay"), []byte("/*"), []byte("*/"), []byte("--"), []byte("@@"),
}
kwXSS = [][]byte{
[]byte("<script"), []byte("</script"), []byte("javascript:"), []byte("onerror="),
[]byte("onload="), []byte("onmouseover="), []byte("alert("), []byte("prompt("),
[]byte("confirm("), []byte("document.cookie"), []byte("document.write"),
[]byte("eval("), []byte("settimeout("), []byte("fromcharcode"), []byte("<iframe"),
[]byte("<svg"), []byte("<img"), []byte("expression("), []byte("vbscript:"),
[]byte("base64,"),
}
kwCMD = [][]byte{
[]byte("/etc/passwd"), []byte("/etc/shadow"), []byte("/bin/sh"), []byte("/bin/bash"),
[]byte("cmd.exe"), []byte("powershell"), []byte("wget "), []byte("curl "),
[]byte("chmod "), []byte("nc -e"), []byte("bash -i"), []byte("whoami"),
[]byte("ipconfig"), []byte("ifconfig"), []byte("&&"), []byte("||"), []byte("$("),
[]byte("`"), []byte("ping -c"), []byte("net user"), []byte("system("),
[]byte("passthru("), []byte("shell_exec("), []byte("popen("),
}
kwTraversal = [][]byte{
[]byte("../"), []byte("..\\"), []byte("%2e%2e"), []byte("..%2f"), []byte("..%5c"),
[]byte("/etc/"), []byte("c:\\"), []byte("c:/"), []byte("web.config"),
[]byte("boot.ini"), []byte("win.ini"), []byte("/proc/self"), []byte("wp-config"),
}
kwProto = [][]byte{
[]byte("<?php"), []byte("<%"), []byte("${"), []byte("#{"), []byte("{{"),
[]byte("jndi:"), []byte("ldap://"), []byte("rmi://"), []byte("file://"),
[]byte("php://"), []byte("data://"), []byte("gopher://"), []byte("dict://"),
[]byte("expect://"), []byte("phar://"),
}
kwScannerUA = [][]byte{
[]byte("sqlmap"), []byte("nikto"), []byte("nmap"), []byte("acunetix"),
[]byte("nessus"), []byte("burp"), []byte("dirbuster"), []byte("wfuzz"),
[]byte("fuzz"), []byte("masscan"), []byte("zgrab"), []byte("python-requests"),
[]byte("go-http-client"), []byte("curl/"), []byte("wget/"), []byte("libwww"),
}
)
// FeatureNames 返回特征名顺序（与 Python feature_names() 一致），用于调试/导出。
func FeatureNames() []string {
names := []string{
"path_len", "query_len", "body_len", "ua_len",
"param_count", "max_param_len", "path_depth", "method_idx",
"special_ratio", "digit_ratio", "letter_ratio", "upper_ratio",
"nonascii_ratio", "entropy", "decode_layers", "query_eq_count",
"kw_sql", "kw_xss", "kw_cmd", "kw_traversal", "kw_proto", "kw_scanner_ua",
}
for i := 0; i < ngramBuckets; i++ {
names = append(names, "ngram_"+itoa(i))
}
return names
}
// 关键词家族在特征向量中的下标（与 ExtractFeatures 布局一致）
const (
idxKwSQL = 16
idxKwXSS = 17
idxKwCMD = 18
idxKwTraversal = 19
idxKwProto = 20
)
// CategoryHint 依据特征里的关键词家族计数，给 AI 命中一个粗粒度、可读且稳定的
// 类别标签，用于日志展示与"按规则汇总"统计（取值是有限小集合，不会像分数那样发散）。
//
// 注意：这只是基于可疑关键词的启发式标注，模型本身只输出"异常概率"，并不真正分类；
// 该函数是 Go 侧的展示辅助，不参与打分、也不属于特征规范的双语言一致性约束。
func CategoryHint(features []float64) string {
if len(features) < scalarFeatureCount {
return "异常请求"
}
cats := []struct {
idx int
name string
}{
{idxKwSQL, "SQL注入"},
{idxKwXSS, "XSS"},
{idxKwCMD, "命令执行"},
{idxKwTraversal, "目录穿越"},
{idxKwProto, "注入攻击"},
}
best := 0.0
bestName := ""
for _, c := range cats {
if features[c.idx] > best { // 严格大于：并列时保留更高优先级（靠前）的类别
best = features[c.idx]
bestName = c.name
}
}
if bestName == "" {
return "异常请求"
}
return bestName
}
func itoa(i int) string {
if i == 0 {
return "0"
}
var buf [4]byte
pos := len(buf)
for i > 0 {
pos--
buf[pos] = byte('0' + i%10)
i /= 10
}
return string(buf[pos:])
}
func truncate(s string) []byte {
b := []byte(s)
if len(b) > maxComponentBytes {
b = b[:maxComponentBytes]
}
return b
}
func isHex(c byte) bool {
return (c >= '0' && c <= '9') || (c >= 'A' && c <= 'F') || (c >= 'a' && c <= 'f')
}
func hexVal(c byte) byte {
switch {
case c >= '0' && c <= '9':
return c - '0'
case c >= 'A' && c <= 'F':
return c - 'A' + 10
default:
return c - 'a' + 10
}
}
// percentDecodeOnce 单次百分号解码：%XX -> 字节；'+' -> 空格；非法 % 序列原样保留。
func percentDecodeOnce(b []byte) []byte {
out := make([]byte, 0, len(b))
n := len(b)
for i := 0; i < n; {
c := b[i]
if c == '%' && i+2 < n && isHex(b[i+1]) && isHex(b[i+2]) {
out = append(out, hexVal(b[i+1])<<4|hexVal(b[i+2]))
i += 3
} else if c == '+' {
out = append(out, ' ')
i++
} else {
out = append(out, c)
i++
}
}
return out
}
// iterDecode 迭代解码至多 decodeMax 次，返回 (解码结果, 实际生效的解码层数)。
func iterDecode(b []byte) ([]byte, int) {
layers := 0
cur := b
for k := 0; k < decodeMax; k++ {
dec := percentDecodeOnce(cur)
if bytes.Equal(dec, cur) {
break
}
cur = dec
layers++
}
return cur, layers
}
// asciiLower 仅对 ASCII 'A'-'Z' 转小写（不处理多字节字符）。
func asciiLower(b []byte) []byte {
out := make([]byte, len(b))
for i, c := range b {
if c >= 'A' && c <= 'Z' {
out[i] = c + 32
} else {
out[i] = c
}
}
return out
}
func fnv1a32(b []byte) uint32 {
var h uint32 = 2166136261
for _, c := range b {
h ^= uint32(c)
h *= 16777619
}
return h
}
func entropy(b []byte) float64 {
if len(b) == 0 {
return 0
}
var counts [256]int
for _, c := range b {
counts[c]++
}
n := float64(len(b))
ent := 0.0
for _, cnt := range counts {
if cnt > 0 {
p := float64(cnt) / n
ent -= p * math.Log2(p)
}
}
return ent
}
func isPunct(c byte) bool {
return (c >= 0x21 && c <= 0x2F) || (c >= 0x3A && c <= 0x40) ||
(c >= 0x5B && c <= 0x60) || (c >= 0x7B && c <= 0x7E)
}
func countKeywords(haystack []byte, kws [][]byte) float64 {
total := 0
for _, kw := range kws {
total += bytes.Count(haystack, kw)
}
return float64(total)
}
// ExtractFeatures 提取特征向量（FeatureCount 维，顺序固定，与 Python 一致）。
//
// 入参均为原始字符串（path/query/body 未解码；query 不含 '?'；method 任意大小写）。
func ExtractFeatures(method, path, query, body, userAgent string) []float64 {
pathB := truncate(path)
queryB := truncate(query)
bodyB := truncate(body)
uaB := truncate(userAgent)
pathD, l1 := iterDecode(pathB)
queryD, l2 := iterDecode(queryB)
bodyD, l3 := iterDecode(bodyB)
decodeLayers := l1
if l2 > decodeLayers {
decodeLayers = l2
}
if l3 > decodeLayers {
decodeLayers = l3
}
combined := make([]byte, 0, len(pathD)+len(queryD)+len(bodyD)+2)
combined = append(combined, pathD...)
combined = append(combined, '\n')
combined = append(combined, queryD...)
combined = append(combined, '\n')
combined = append(combined, bodyD...)
combinedLower := asciiLower(combined)
uaLower := asciiLower(uaB)
f := make([]float64, FeatureCount)
// --- 长度/结构类 ---
f[0] = float64(len(pathB))
f[1] = float64(len(queryB))
f[2] = float64(len(bodyB))
f[3] = float64(len(uaB))
paramCount := 0
maxParamLen := 0
for _, seg := range bytes.Split(queryD, []byte("&")) {
if len(seg) == 0 {
continue
}
paramCount++
eq := bytes.IndexByte(seg, '=')
vlen := 0
if eq >= 0 {
vlen = len(seg) - eq - 1
}
if vlen > maxParamLen {
maxParamLen = vlen
}
}
f[4] = float64(paramCount)
f[5] = float64(maxParamLen)
f[6] = float64(bytes.Count(pathD, []byte("/")))
if v, ok := methodIdx[upperASCII(method)]; ok {
f[7] = v
} else {
f[7] = methodIdxOther
}
// --- 字符分布类（在 combined 上统计）---
n := len(combined)
if n > 0 {
var punct, digit, letter, upper, nonascii int
for _, c := range combined {
if isPunct(c) {
punct++
}
if c >= '0' && c <= '9' {
digit++
}
if (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') {
letter++
}
if c >= 'A' && c <= 'Z' {
upper++
}
if c >= 0x80 {
nonascii++
}
}
fn := float64(n)
f[8] = float64(punct) / fn
f[9] = float64(digit) / fn
f[10] = float64(letter) / fn
f[11] = float64(upper) / fn
f[12] = float64(nonascii) / fn
}
f[13] = entropy(combined)
f[14] = float64(decodeLayers)
f[15] = float64(bytes.Count(queryD, []byte("=")))
// --- 关键词类 ---
f[16] = countKeywords(combinedLower, kwSQL)
f[17] = countKeywords(combinedLower, kwXSS)
f[18] = countKeywords(combinedLower, kwCMD)
f[19] = countKeywords(combinedLower, kwTraversal)
f[20] = countKeywords(combinedLower, kwProto)
f[21] = countKeywords(uaLower, kwScannerUA)
// --- 字节 3-gram 哈希桶频率 ---
ln := len(combinedLower)
if ln >= ngramN {
var counts [ngramBuckets]int
totalNgrams := ln - ngramN + 1
for i := 0; i < totalNgrams; i++ {
h := fnv1a32(combinedLower[i : i+ngramN])
counts[h%ngramBuckets]++
}
for bi := 0; bi < ngramBuckets; bi++ {
f[scalarFeatureCount+bi] = float64(counts[bi]) / float64(totalNgrams)
}
}
return f
}
func upperASCII(s string) string {
b := []byte(s)
changed := false
for i, c := range b {
if c >= 'a' && c <= 'z' {
b[i] = c - 32
changed = true
}
}
if !changed {
return s
}
return string(b)
}
wafai/features_test.go
+115
View File
@@ -0,0 +1,115 @@
package wafai
import (
"encoding/json"
"math"
"os"
"path/filepath"
"testing"
)
// goldenFile 由 SamWafAI 生成：uv run samwafai gen-golden --out wafai/testdata/feature_golden.json
type goldenFile struct {
FeatureVersion string `json:"feature_version"`
FeatureNames []string `json:"feature_names"`
Cases []struct {
Method string `json:"method"`
Path string `json:"path"`
Query string `json:"query"`
Body string `json:"body"`
UserAgent string `json:"user_agent"`
Features []float64 `json:"features"`
} `json:"cases"`
}
// TestFeatureParity 校验 Go 特征提取与 Python（golden 文件）逐项一致。
// 这是 AI 检测正确性的基石：特征不对齐 -> 模型打分错乱。
func TestFeatureParity(t *testing.T) {
raw, err := os.ReadFile(filepath.Join("testdata", "feature_golden.json"))
if err != nil {
t.Fatalf("读取 golden 文件失败（请先运行 uv run samwafai gen-golden）: %v", err)
}
var g goldenFile
if err := json.Unmarshal(raw, &g); err != nil {
t.Fatalf("解析 golden 文件失败: %v", err)
}
if g.FeatureVersion != FeatureVersion {
t.Fatalf("特征版本不一致: golden=%s go=%s", g.FeatureVersion, FeatureVersion)
}
if len(g.FeatureNames) != FeatureCount {
t.Fatalf("特征维度不一致: golden=%d go=%d", len(g.FeatureNames), FeatureCount)
}
goNames := FeatureNames()
for i, name := range g.FeatureNames {
if goNames[i] != name {
t.Fatalf("特征名[%d]不一致: golden=%s go=%s", i, name, goNames[i])
}
}
const eps = 1e-9
for ci, c := range g.Cases {
got := ExtractFeatures(c.Method, c.Path, c.Query, c.Body, c.UserAgent)
if len(got) != len(c.Features) {
t.Fatalf("case[%d] 维度不一致: got=%d want=%d", ci, len(got), len(c.Features))
}
for fi := range got {
if math.Abs(got[fi]-c.Features[fi]) > eps {
t.Errorf("case[%d] 特征[%d:%s] 不一致: go=%.12f py=%.12f (path=%q query=%q)",
ci, fi, goNames[fi], got[fi], c.Features[fi], c.Path, c.Query)
}
}
}
}
func TestPercentDecode(t *testing.T) {
cases := []struct {
in string
want string
}{
{"%2e%2e", ".."},
{"a+b", "a b"},
{"%zz", "%zz"},
{"%2", "%2"},
}
for _, c := range cases {
if got := string(percentDecodeOnce([]byte(c.in))); got != c.want {
t.Errorf("percentDecodeOnce(%q)=%q want %q", c.in, got, c.want)
}
}
}
func TestIterDecodeLayers(t *testing.T) {
dec, layers := iterDecode([]byte("%252e%252e"))
if string(dec) != ".." || layers != 2 {
t.Errorf("iterDecode 双层解码失败: got=%q layers=%d", string(dec), layers)
}
}
func TestCategoryHint(t *testing.T) {
cases := []struct {
method, path, query, body, ua string
want string
}{
{"GET", "/p", "id=1' or '1'='1 union select", "", "", "SQL注入"},
{"GET", "/c", "q=<script>alert(1)</script>", "", "", "XSS"},
{"GET", "/d", "file=../../../../etc/passwd", "", "", "目录穿越"},
{"GET", "/x", "host=127.0.0.1;cat /etc/passwd", "", "", "命令执行"},
{"GET", "/", "page=1&sort=asc", "", "Mozilla/5.0", "异常请求"},
}
for _, c := range cases {
f := ExtractFeatures(c.method, c.path, c.query, c.body, c.ua)
if got := CategoryHint(f); got != c.want {
t.Errorf("CategoryHint(query=%q)=%q want %q", c.query, got, c.want)
}
}
}
func TestFNV1a(t *testing.T) {
if fnv1a32([]byte("")) != 2166136261 {
t.Error("fnv1a32 空串错误")
}
if fnv1a32([]byte("a")) != 0xE40C292C {
t.Errorf("fnv1a32('a')=%#x", fnv1a32([]byte("a")))
}
}
wafai/labeling.go
+102
View File
@@ -0,0 +1,102 @@
package wafai
import "strings"
// 规则弱标签：把 SamWaf 日志里的规则命中转换为训练标签。
// 与 SamWafAI samwafai/data/labeling.py 保持一致的判定逻辑。
//
// 只把"载荷型"攻击规则的命中当作正样本；CC/IP黑白名单/防盗链/验证码等
// 行为型或策略型拦截与请求载荷内容无关，全部丢弃，避免污染训练。
// payloadRuleKeywords RULE 标题包含这些关键字 -> 载荷型攻击（正样本）
var payloadRuleKeywords = []struct {
kw string
attackType string
}{
{"sql", "sqli"}, {"注入", "sqli"},
{"xss", "xss"}, {"跨站", "xss"},
{"rce", "rce"}, {"命令执行", "rce"}, {"代码执行", "rce"},
{"扫描", "scan"}, {"scan", "scan"},
{"穿越", "traversal"}, {"traversal", "traversal"},
{"owasp", "owasp"},
}
// excludeRuleKeywords RULE 标题包含这些关键字 -> 非载荷型，直接丢弃
var excludeRuleKeywords = []string{
"cc", "频次", "rate limit",
"ip", "黑名单", "白名单",
"防盗链", "盗链",
"验证码", "captcha",
"敏感词",
"bot", "爬虫", "蜘蛛",
"url黑", "url白", "禁止url", "路径禁止",
"应用",
"ai检测", // AI 自己产生的命中不能再回灌当弱标签，避免自我强化
}
var blockActions = []string{"禁止", "阻止"}
// highConfidenceAttackTypes 高置信攻击类型：命中即基本可判定为真攻击（误报率低），
// 未经人工确认也可作为训练正样本。
// - sqli/xss：libinjection 词法分析，精度高。
// - owasp：OWASP CRS 规则集。
// - rce：检测仅匹配 phpinfo()/call_user_func_array/invokefunction 等极具体签名，正常流量几乎不出现，误报极低。
// - traversal：匹配 ../、..\、%2e%2e 等穿越特征，正常流量很少出现，误报低。
//
// scan 由 User-Agent 判定（curl/python-requests 等合法工具也会命中），误报高，仍排除在高置信外，
// 未经人工确认不当攻击；自定义规则同理。
var highConfidenceAttackTypes = map[string]bool{
"sqli": true,
"xss": true,
"owasp": true,
"rce": true,
"traversal": true,
}
// IsHighConfidenceAttackType 该攻击类型是否高置信（可在无人工确认时信任为正样本）。
func IsHighConfidenceAttackType(attackType string) bool {
return highConfidenceAttackTypes[attackType]
}
// WeakLabelVerdict 弱标签判定结果。
type WeakLabelVerdict int
const (
VerdictDrop WeakLabelVerdict = iota // 丢弃（不用于训练）
VerdictAttack // 载荷型攻击正样本
VerdictNormal // 正常负样本
)
// WeakLabel 根据日志的 action/rule/logOnlyMode 给出弱标签。
func WeakLabel(action, rule string, logOnlyMode int) (WeakLabelVerdict, string) {
ruleL := strings.ToLower(strings.TrimSpace(rule))
blocked := false
for _, a := range blockActions {
if strings.TrimSpace(action) == a {
blocked = true
break
}
}
if logOnlyMode == 1 && ruleL != "" {
blocked = true
}
if !blocked {
if ruleL == "" {
return VerdictNormal, ""
}
return VerdictDrop, ""
}
for _, p := range payloadRuleKeywords {
if strings.Contains(ruleL, p.kw) {
return VerdictAttack, p.attackType
}
}
for _, kw := range excludeRuleKeywords {
if strings.Contains(ruleL, kw) {
return VerdictDrop, ""
}
}
return VerdictDrop, ""
}
wafai/manager.go
+120
View File
@@ -0,0 +1,120 @@
package wafai
import (
"fmt"
"sync"
"sync/atomic"
)
// Detector AI 检测器：持有当前模型，提供线程安全的打分与热加载。
//
// 失败安全：模型未加载/打分 panic 时，Predict 返回非命中结果，绝不影响业务转发。
type Detector struct {
mu sync.RWMutex
current atomic.Pointer[loadedModel]
}
type loadedModel struct {
engine InferenceEngine
manifest Manifest
}
// PredictResult 单次推理结果。
type PredictResult struct {
Loaded bool // 当前是否有可用模型
Score float64 // 攻击概率 [0,1]
Category string // 粗粒度类别提示（SQL注入/XSS/.../异常请求），用于日志展示与按规则汇总
BlockThreshold float64 // 模型建议的拦截阈值
ObserveThreshold float64 // 模型建议的观察阈值
ModelVersion string
}
// NewDetector 创建空检测器（尚未加载模型）。
func NewDetector() *Detector {
return &Detector{}
}
// LoadFromFile 从 .swai 文件加载模型并原子热替换。
func (d *Detector) LoadFromFile(swaiPath string) (Manifest, error) {
pkg, err := loadPackageFile(swaiPath)
if err != nil {
return Manifest{}, err
}
return d.loadPackage(pkg)
}
// LoadFromBytes 从内存字节加载模型并原子热替换。
func (d *Detector) LoadFromBytes(data []byte) (Manifest, error) {
pkg, err := loadPackageBytes(data)
if err != nil {
return Manifest{}, err
}
return d.loadPackage(pkg)
}
func (d *Detector) loadPackage(pkg *loadedPackage) (Manifest, error) {
d.mu.Lock()
defer d.mu.Unlock()
var engine InferenceEngine
var err error
switch pkg.Manifest.ModelType {
case "gbdt", "":
engine, err = newGBDTEngine(pkg.ModelBytes, pkg.Manifest.FeatureVersion)
case "onnx":
engine, err = newONNXEngine(pkg.ModelBytes, pkg.Manifest.FeatureVersion)
default:
return Manifest{}, fmt.Errorf("不支持的模型类型: %q", pkg.Manifest.ModelType)
}
if err != nil {
return Manifest{}, err
}
d.current.Store(&loadedModel{engine: engine, manifest: pkg.Manifest})
return pkg.Manifest, nil
}
// Unload 卸载当前模型。
func (d *Detector) Unload() {
d.current.Store(nil)
}
// IsLoaded 当前是否已加载模型。
func (d *Detector) IsLoaded() bool {
return d.current.Load() != nil
}
// CurrentManifest 返回当前模型 manifest（未加载时第二个返回值为 false）。
func (d *Detector) CurrentManifest() (Manifest, bool) {
m := d.current.Load()
if m == nil {
return Manifest{}, false
}
return m.manifest, true
}
// PredictRequest 对一次请求的关键字段打分。失败安全：任何异常返回 Loaded=false。
func (d *Detector) PredictRequest(method, path, query, body, userAgent string) (res PredictResult) {
lm := d.current.Load()
if lm == nil {
return PredictResult{Loaded: false}
}
defer func() {
if r := recover(); r != nil {
// 推理 panic 时降级为未命中，绝不阻断业务
res = PredictResult{Loaded: false}
}
}()
features := ExtractFeatures(method, path, query, body, userAgent)
score := lm.engine.Score(features)
return PredictResult{
Loaded: true,
Score: score,
Category: CategoryHint(features),
BlockThreshold: lm.manifest.BlockThreshold,
ObserveThreshold: lm.manifest.ObserveThreshold,
ModelVersion: lm.manifest.ModelVersion,
}
}
wafai/modelpack.go
+150
View File
@@ -0,0 +1,150 @@
package wafai
import (
"archive/zip"
"bytes"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"path"
)
const (
swaiFormatVersion = 1
manifestName = "manifest.json"
maxModelBytes = 64 * 1024 * 1024 // 单个模型文件解压上限 64MB（防 zip 炸弹）
maxManifestBytes = 1 * 1024 * 1024 // manifest 上限 1MB
maxPackageEntries = 16 // 包内条目数上限
)
// Manifest .swai 模型包元数据（与 SamWafAI export/package.py 对齐）。
type Manifest struct {
SwaiFormatVersion int `json:"swai_format_version"`
ModelVersion string `json:"model_version"`
ModelType string `json:"model_type"`
FeatureVersion string `json:"feature_version"`
FeatureCount int `json:"feature_count"`
ModelFile string `json:"model_file"`
ModelSha256 string `json:"model_sha256"`
BlockThreshold float64 `json:"block_threshold"`
ObserveThreshold float64 `json:"observe_threshold"`
DataFingerprint string `json:"data_fingerprint"`
CreatedAt string `json:"created_at"`
}
// loadedPackage 解析后的模型包内容。
type loadedPackage struct {
Manifest Manifest
ModelBytes []byte
}
// loadPackageFile 从磁盘加载并校验 .swai 模型包。
func loadPackageFile(swaiPath string) (*loadedPackage, error) {
zr, err := zip.OpenReader(swaiPath)
if err != nil {
return nil, fmt.Errorf("打开模型包失败: %w", err)
}
defer zr.Close()
return parsePackage(&zr.Reader)
}
// loadPackageBytes 从内存字节加载并校验 .swai 模型包。
func loadPackageBytes(data []byte) (*loadedPackage, error) {
zr, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
if err != nil {
return nil, fmt.Errorf("解析模型包失败: %w", err)
}
return parsePackage(zr)
}
func parsePackage(zr *zip.Reader) (*loadedPackage, error) {
if len(zr.File) > maxPackageEntries {
return nil, fmt.Errorf("模型包条目过多: %d", len(zr.File))
}
var manifestRaw []byte
files := map[string][]byte{}
for _, zf := range zr.File {
// 防路径穿越：条目名必须是简单文件名
name := zf.Name
if path.IsAbs(name) || name != path.Clean(name) || containsDotDot(name) {
return nil, fmt.Errorf("非法的模型包条目名: %q", name)
}
limit := int64(maxModelBytes)
if name == manifestName {
limit = maxManifestBytes
}
data, err := readZipEntry(zf, limit)
if err != nil {
return nil, fmt.Errorf("读取条目 %q 失败: %w", name, err)
}
if name == manifestName {
manifestRaw = data
} else {
files[name] = data
}
}
if manifestRaw == nil {
return nil, errors.New("模型包缺少 manifest.json")
}
var m Manifest
if err := json.Unmarshal(manifestRaw, &m); err != nil {
return nil, fmt.Errorf("解析 manifest 失败: %w", err)
}
if m.SwaiFormatVersion != swaiFormatVersion {
return nil, fmt.Errorf("不支持的模型包格式版本: %d（当前支持 %d）", m.SwaiFormatVersion, swaiFormatVersion)
}
// 特征版本硬约束：与 Go 侧不一致则拒绝加载
if m.FeatureVersion != FeatureVersion {
return nil, fmt.Errorf("特征版本不匹配: 模型=%s, 引擎=%s（请用匹配版本的 SamWafAI 重新训练）", m.FeatureVersion, FeatureVersion)
}
if m.FeatureCount != 0 && m.FeatureCount != FeatureCount {
return nil, fmt.Errorf("特征维度不匹配: 模型=%d, 引擎=%d", m.FeatureCount, FeatureCount)
}
modelBytes, ok := files[m.ModelFile]
if !ok {
return nil, fmt.Errorf("模型包缺少模型文件: %q", m.ModelFile)
}
// sha256 完整性校验
if m.ModelSha256 != "" {
sum := sha256.Sum256(modelBytes)
if got := hex.EncodeToString(sum[:]); got != m.ModelSha256 {
return nil, fmt.Errorf("模型文件 sha256 校验失败: 期望 %s 实际 %s", m.ModelSha256, got)
}
}
return &loadedPackage{Manifest: m, ModelBytes: modelBytes}, nil
}
func readZipEntry(zf *zip.File, limit int64) ([]byte, error) {
rc, err := zf.Open()
if err != nil {
return nil, err
}
defer rc.Close()
// 限制读取量防 zip 炸弹：多读 1 字节用于判断是否超限
data, err := io.ReadAll(io.LimitReader(rc, limit+1))
if err != nil {
return nil, err
}
if int64(len(data)) > limit {
return nil, fmt.Errorf("条目解压大小超过上限 %d 字节", limit)
}
return data, nil
}
func containsDotDot(p string) bool {
for i := 0; i+1 < len(p); i++ {
if p[i] == '.' && p[i+1] == '.' {
return true
}
}
return false
}
wafai/modelpack_test.go
+68
View File
@@ -0,0 +1,68 @@
package wafai
import (
"os"
"path/filepath"
"testing"
)
// TestLoadAndScore 端到端：加载 SamWafAI 产出的 .swai 模型，对正常/攻击请求打分。
// 模型由合成数据训练（examples/gen_synthetic.py），攻击样本应显著高于正常样本。
func TestLoadAndScore(t *testing.T) {
swaiPath := filepath.Join("testdata", "sample_model.swai")
if _, err := os.Stat(swaiPath); err != nil {
t.Skip("无 sample_model.swai，跳过（由 SamWafAI pipeline 生成）")
}
d := NewDetector()
if d.IsLoaded() {
t.Fatal("新建检测器不应已加载")
}
// 未加载时打分应失败安全
if res := d.PredictRequest("GET", "/", "id=1", "", ""); res.Loaded {
t.Fatal("未加载模型时 Loaded 应为 false")
}
manifest, err := d.LoadFromFile(swaiPath)
if err != nil {
t.Fatalf("加载模型失败: %v", err)
}
if manifest.FeatureVersion != FeatureVersion {
t.Fatalf("特征版本不一致: %s", manifest.FeatureVersion)
}
if !d.IsLoaded() {
t.Fatal("加载后 IsLoaded 应为 true")
}
normal := d.PredictRequest("GET", "/products", "page=1&sort=asc", "", "Mozilla/5.0")
attack := d.PredictRequest("GET", "/index.php", "id=1' or '1'='1 union select null,null--", "", "sqlmap/1.5")
if !normal.Loaded || !attack.Loaded {
t.Fatal("加载后打分应成功")
}
t.Logf("normal score=%.4f attack score=%.4f block_thr=%.4f", normal.Score, attack.Score, manifest.BlockThreshold)
if attack.Score <= normal.Score {
t.Errorf("攻击样本分数(%.4f)应高于正常样本(%.4f)", attack.Score, normal.Score)
}
d.Unload()
if d.IsLoaded() {
t.Fatal("卸载后 IsLoaded 应为 false")
}
}
func TestRejectBadFeatureVersion(t *testing.T) {
// 篡改 manifest 的特征版本应被拒绝（构造最小 zip）
bad := buildTestPackage(t, `{"swai_format_version":1,"feature_version":"v999","model_type":"gbdt","model_file":"model_lgbm.txt","model_sha256":""}`, "dummy")
d := NewDetector()
if _, err := d.LoadFromBytes(bad); err == nil {
t.Fatal("特征版本不匹配应加载失败")
}
}
func TestRejectPathTraversalEntry(t *testing.T) {
bad := buildTestPackageNamed(t, manifestName, `{"swai_format_version":1,"feature_version":"v1"}`, "../evil.txt", "x")
d := NewDetector()
if _, err := d.LoadFromBytes(bad); err == nil {
t.Fatal("含路径穿越条目应加载失败")
}
}
wafai/testdata/feature_golden.json
Vendored
+853
View File
@@ -0,0 +1,853 @@
{
"feature_version": "v1",
"feature_names": [
"path_len",
"query_len",
"body_len",
"ua_len",
"param_count",
"max_param_len",
"path_depth",
"method_idx",
"special_ratio",
"digit_ratio",
"letter_ratio",
"upper_ratio",
"nonascii_ratio",
"entropy",
"decode_layers",
"query_eq_count",
"kw_sql",
"kw_xss",
"kw_cmd",
"kw_traversal",
"kw_proto",
"kw_scanner_ua",
"ngram_0",
"ngram_1",
"ngram_2",
"ngram_3",
"ngram_4",
"ngram_5",
"ngram_6",
"ngram_7",
"ngram_8",
"ngram_9",
"ngram_10",
"ngram_11",
"ngram_12",
"ngram_13",
"ngram_14",
"ngram_15",
"ngram_16",
"ngram_17",
"ngram_18",
"ngram_19",
"ngram_20",
"ngram_21",
"ngram_22",
"ngram_23",
"ngram_24",
"ngram_25",
"ngram_26",
"ngram_27",
"ngram_28",
"ngram_29",
"ngram_30",
"ngram_31",
"ngram_32",
"ngram_33",
"ngram_34",
"ngram_35",
"ngram_36",
"ngram_37",
"ngram_38",
"ngram_39",
"ngram_40",
"ngram_41",
"ngram_42",
"ngram_43",
"ngram_44",
"ngram_45",
"ngram_46",
"ngram_47",
"ngram_48",
"ngram_49",
"ngram_50",
"ngram_51",
"ngram_52",
"ngram_53",
"ngram_54",
"ngram_55",
"ngram_56",
"ngram_57",
"ngram_58",
"ngram_59",
"ngram_60",
"ngram_61",
"ngram_62",
"ngram_63"
],
"cases": [
{
"method": "GET",
"path": "/",
"query": "",
"body": "",
"user_agent": "",
"features": [
1.0,
0.0,
0.0,
0.0,
0.0,
0.0,
1.0,
0.0,
0.3333333333333333,
0.0,
0.0,
0.0,
0.0,
0.9182958340544896,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
1.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0
]
},
{
"method": "GET",
"path": "/index.php",
"query": "id=1' or '1'='1",
"body": "",
"user_agent": "Mozilla/5.0",
"features": [
10.0,
15.0,
0.0,
11.0,
1.0,
12.0,
1.0,
0.0,
0.2962962962962963,
0.1111111111111111,
0.4444444444444444,
0.0,
0.0,
3.838039816898156,
0.0,
2.0,
1.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.04,
0.0,
0.0,
0.04,
0.0,
0.0,
0.0,
0.0,
0.0,
0.04,
0.04,
0.0,
0.0,
0.0,
0.04,
0.04,
0.04,
0.0,
0.0,
0.0,
0.04,
0.0,
0.0,
0.0,
0.0,
0.0,
0.04,
0.0,
0.0,
0.16,
0.04,
0.0,
0.0,
0.0,
0.04,
0.04,
0.04,
0.0,
0.0,
0.0,
0.04,
0.0,
0.04,
0.04,
0.0,
0.0,
0.04,
0.0,
0.04,
0.0,
0.04,
0.0,
0.0,
0.0,
0.0,
0.04,
0.0,
0.04,
0.0,
0.0
]
},
{
"method": "POST",
"path": "/login",
"query": "",
"body": "user=admin&pwd=123",
"user_agent": "sqlmap/1.5",
"features": [
6.0,
0.0,
18.0,
10.0,
0.0,
0.0,
1.0,
1.0,
0.15384615384615385,
0.11538461538461539,
0.6538461538461539,
0.0,
0.0,
4.315824333525707,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
1.0,
0.041666666666666664,
0.0,
0.0,
0.0,
0.041666666666666664,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.041666666666666664,
0.0,
0.0,
0.0,
0.0,
0.0,
0.08333333333333333,
0.0,
0.0,
0.0,
0.0,
0.041666666666666664,
0.041666666666666664,
0.0,
0.0,
0.0,
0.125,
0.0,
0.041666666666666664,
0.0,
0.041666666666666664,
0.0,
0.041666666666666664,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.041666666666666664,
0.08333333333333333,
0.0,
0.0,
0.0,
0.0,
0.0,
0.041666666666666664,
0.0,
0.125,
0.0,
0.0,
0.08333333333333333,
0.041666666666666664,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.041666666666666664,
0.0,
0.0,
0.0
]
},
{
"method": "GET",
"path": "/a/b/c",
"query": "q=%3Cscript%3Ealert(1)%3C%2Fscript%3E",
"body": "",
"user_agent": "curl/7.0",
"features": [
6.0,
37.0,
0.0,
8.0,
1.0,
25.0,
3.0,
0.0,
0.3142857142857143,
0.02857142857142857,
0.6,
0.0,
0.0,
4.09314980247381,
1.0,
1.0,
0.0,
3.0,
0.0,
0.0,
0.0,
1.0,
0.0,
0.09090909090909091,
0.0,
0.0,
0.030303030303030304,
0.030303030303030304,
0.030303030303030304,
0.0,
0.030303030303030304,
0.0,
0.0,
0.0,
0.06060606060606061,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.06060606060606061,
0.0,
0.0,
0.0,
0.030303030303030304,
0.0,
0.0,
0.0,
0.0,
0.030303030303030304,
0.09090909090909091,
0.0,
0.0,
0.030303030303030304,
0.06060606060606061,
0.0,
0.0,
0.0,
0.0,
0.030303030303030304,
0.030303030303030304,
0.0,
0.030303030303030304,
0.0,
0.06060606060606061,
0.0,
0.06060606060606061,
0.0,
0.0,
0.0,
0.0,
0.030303030303030304,
0.0,
0.0,
0.0,
0.0,
0.030303030303030304,
0.06060606060606061,
0.0,
0.0,
0.06060606060606061,
0.0,
0.0,
0.0,
0.030303030303030304
]
},
{
"method": "GET",
"path": "/download",
"query": "file=../../../../etc/passwd",
"body": "",
"user_agent": "",
"features": [
9.0,
27.0,
0.0,
0.0,
1.0,
22.0,
1.0,
0.0,
0.39473684210526316,
0.0,
0.5526315789473685,
0.0,
0.0,
3.7146469211675206,
0.0,
1.0,
0.0,
0.0,
1.0,
5.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.1111111111111111,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.05555555555555555,
0.0,
0.027777777777777776,
0.0,
0.05555555555555555,
0.0,
0.0,
0.05555555555555555,
0.027777777777777776,
0.027777777777777776,
0.0,
0.05555555555555555,
0.027777777777777776,
0.1111111111111111,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.05555555555555555,
0.0,
0.0,
0.0,
0.05555555555555555,
0.027777777777777776,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.027777777777777776,
0.027777777777777776,
0.027777777777777776,
0.0,
0.027777777777777776,
0.0,
0.0,
0.0,
0.027777777777777776,
0.1111111111111111,
0.027777777777777776,
0.0,
0.027777777777777776
]
},
{
"method": "GET",
"path": "/cmd",
"query": "x=%2e%2e%252fetc%252fpasswd",
"body": "",
"user_agent": "",
"features": [
4.0,
27.0,
0.0,
0.0,
1.0,
13.0,
1.0,
0.0,
0.2857142857142857,
0.0,
0.6190476190476191,
0.0,
0.0,
3.689703732199547,
2.0,
1.0,
0.0,
0.0,
1.0,
2.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.05263157894736842,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.10526315789473684,
0.0,
0.05263157894736842,
0.0,
0.10526315789473684,
0.0,
0.0,
0.10526315789473684,
0.05263157894736842,
0.0,
0.0,
0.0,
0.0,
0.05263157894736842,
0.0,
0.0,
0.0,
0.05263157894736842,
0.0,
0.05263157894736842,
0.0,
0.0,
0.0,
0.0,
0.0,
0.05263157894736842,
0.0,
0.0,
0.0,
0.0,
0.05263157894736842,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.05263157894736842,
0.05263157894736842,
0.05263157894736842,
0.05263157894736842,
0.05263157894736842,
0.0,
0.0
]
},
{
"method": "PUT",
"path": "/upload/中文路径",
"query": "name=测试&v=1",
"body": "数据body",
"user_agent": "Go-http-client/1.1",
"features": [
20.0,
15.0,
10.0,
18.0,
2.0,
6.0,
2.0,
2.0,
0.10638297872340426,
0.02127659574468085,
0.3191489361702128,
0.0,
0.5106382978723404,
5.001397362315936,
0.0,
2.0,
0.0,
0.0,
0.0,
0.0,
0.0,
1.0,
0.022222222222222223,
0.0,
0.044444444444444446,
0.022222222222222223,
0.022222222222222223,
0.0,
0.0,
0.044444444444444446,
0.0,
0.0,
0.0,
0.044444444444444446,
0.022222222222222223,
0.022222222222222223,
0.0,
0.0,
0.022222222222222223,
0.0,
0.044444444444444446,
0.044444444444444446,
0.0,
0.06666666666666667,
0.022222222222222223,
0.022222222222222223,
0.0,
0.0,
0.044444444444444446,
0.0,
0.022222222222222223,
0.044444444444444446,
0.0,
0.0,
0.022222222222222223,
0.0,
0.022222222222222223,
0.0,
0.022222222222222223,
0.0,
0.022222222222222223,
0.022222222222222223,
0.022222222222222223,
0.0,
0.0,
0.022222222222222223,
0.0,
0.022222222222222223,
0.06666666666666667,
0.044444444444444446,
0.0,
0.0,
0.022222222222222223,
0.0,
0.022222222222222223,
0.0,
0.0,
0.022222222222222223,
0.022222222222222223,
0.0,
0.0,
0.0,
0.044444444444444446,
0.0,
0.0,
0.0
]
},
{
"method": "POST",
"path": "/api",
"query": "",
"body": "{\"q\":\"union select null,null--\"}",
"user_agent": "python-requests/2.0",
"features": [
4.0,
0.0,
32.0,
19.0,
0.0,
0.0,
1.0,
1.0,
0.2894736842105263,
0.0,
0.6052631578947368,
0.0,
0.0,
4.133071514059368,
0.0,
0.0,
3.0,
0.0,
0.0,
0.0,
0.0,
1.0,
0.027777777777777776,
0.0,
0.027777777777777776,
0.0,
0.0,
0.0,
0.0,
0.027777777777777776,
0.05555555555555555,
0.027777777777777776,
0.027777777777777776,
0.0,
0.0,
0.027777777777777776,
0.0,
0.0,
0.08333333333333333,
0.027777777777777776,
0.0,
0.05555555555555555,
0.0,
0.0,
0.0,
0.027777777777777776,
0.027777777777777776,
0.0,
0.0,
0.027777777777777776,
0.0,
0.0,
0.0,
0.0,
0.08333333333333333,
0.0,
0.0,
0.0,
0.0,
0.027777777777777776,
0.0,
0.027777777777777776,
0.027777777777777776,
0.027777777777777776,
0.0,
0.08333333333333333,
0.0,
0.027777777777777776,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.027777777777777776,
0.027777777777777776,
0.027777777777777776,
0.027777777777777776,
0.0,
0.05555555555555555,
0.0,
0.0,
0.027777777777777776,
0.027777777777777776,
0.0
]
}
]
}
wafai/testdata/sample_model.swai
Vendored
BIN
View File
Binary file not shown.
wafai/testhelper_test.go
+33
View File
@@ -0,0 +1,33 @@
package wafai
import (
"archive/zip"
"bytes"
"testing"
)
// buildTestPackage 构造一个含 manifest + 单模型文件的内存 zip（仅供测试）。
func buildTestPackage(t *testing.T, manifestJSON, modelContent string) []byte {
t.Helper()
return buildTestPackageNamed(t, manifestName, manifestJSON, "model_lgbm.txt", modelContent)
}
func buildTestPackageNamed(t *testing.T, manifestEntry, manifestJSON, modelEntry, modelContent string) []byte {
t.Helper()
var buf bytes.Buffer
zw := zip.NewWriter(&buf)
if w, err := zw.Create(manifestEntry); err == nil {
_, _ = w.Write([]byte(manifestJSON))
} else {
t.Fatal(err)
}
if w, err := zw.Create(modelEntry); err == nil {
_, _ = w.Write([]byte(modelContent))
} else {
t.Fatal(err)
}
if err := zw.Close(); err != nil {
t.Fatal(err)
}
return buf.Bytes()
}
wafdb/migrations_core.go
+48
View File
@@ -989,6 +989,54 @@ func RunCoreDBMigrations(db *gorm.DB) error {
return tx.Migrator().DropTable(&model.WafAppChangeLog{})
},
},
{
ID: "202606120003_add_waf_log_label_marks_table",
Migrate: func(tx *gorm.DB) error {
zlog.Info("迁移 202606120003: 创建 waf_log_label_marks 表（AI训练标签人工修正）")
if err := tx.AutoMigrate(&model.WafLogLabelMark{}); err != nil {
return fmt.Errorf("创建 waf_log_label_marks 表失败: %w", err)
}
zlog.Info("waf_log_label_marks 表创建成功")
return nil
},
Rollback: func(tx *gorm.DB) error {
zlog.Info("回滚 202606120003: 删除 waf_log_label_marks 表")
return tx.Migrator().DropTable(&model.WafLogLabelMark{})
},
},
{
ID: "202606120004_add_label_mark_attack_type",
Migrate: func(tx *gorm.DB) error {
zlog.Info("迁移 202606120004: 为 waf_log_label_marks 添加 attack_type 字段")
if tx.Migrator().HasColumn(&model.WafLogLabelMark{}, "attack_type") {
return nil
}
if err := tx.Migrator().AddColumn(&model.WafLogLabelMark{}, "AttackType"); err != nil {
return fmt.Errorf("添加 attack_type 字段失败: %w", err)
}
return nil
},
Rollback: func(tx *gorm.DB) error {
if tx.Migrator().HasColumn(&model.WafLogLabelMark{}, "attack_type") {
return tx.Migrator().DropColumn(&model.WafLogLabelMark{}, "AttackType")
}
return nil
},
},
{
ID: "202606120005_add_label_mark_snapshot_fields",
Migrate: func(tx *gorm.DB) error {
zlog.Info("迁移 202606120005: 为 waf_log_label_marks 添加请求快照字段")
// AutoMigrate 仅新增缺失列，幂等安全
if err := tx.AutoMigrate(&model.WafLogLabelMark{}); err != nil {
return fmt.Errorf("同步 waf_log_label_marks 快照字段失败: %w", err)
}
return nil
},
Rollback: func(tx *gorm.DB) error {
return nil
},
},
})
// 执行迁移
wafdb/migrations_log.go
+40
View File
@@ -188,6 +188,46 @@ func RunLogDBMigrations(db *gorm.DB) error {
return nil
},
},
// 迁移6: 为 web_logs 表添加 ai_score 字段（AI智能检测得分，支持集中查看观察/拦截）
{
ID: "202606120001_add_web_logs_ai_score",
Migrate: func(tx *gorm.DB) error {
zlog.Info("迁移 202606120001: 为 web_logs 表添加 ai_score 字段")
if tx.Migrator().HasColumn(&innerbean.WebLog{}, "ai_score") {
zlog.Info("ai_score 字段已存在，跳过添加")
return nil
}
if err := tx.Migrator().AddColumn(&innerbean.WebLog{}, "AI_SCORE"); err != nil {
return fmt.Errorf("添加 ai_score 字段失败: %w", err)
}
zlog.Info("ai_score 字段添加成功（用于记录AI检测得分，支持观察/拦截集中查看）")
return nil
},
Rollback: func(tx *gorm.DB) error {
zlog.Info("回滚 202606120001: 删除 web_logs 表的 ai_score 字段")
if tx.Migrator().HasColumn(&innerbean.WebLog{}, "ai_score") {
return tx.Migrator().DropColumn(&innerbean.WebLog{}, "AI_SCORE")
}
return nil
},
},
// 迁移7: 为 web_logs 表的 ai_score 建索引（AI看板按 ai_score>0 过滤，命中是极小子集，
// 走索引可避免全表扫描；复合 day 兼顾按天范围过滤与趋势排序）
{
ID: "202606120002_add_web_logs_ai_score_index",
Migrate: func(tx *gorm.DB) error {
zlog.Info("迁移 202606120002: 为 web_logs.ai_score 创建索引")
return safeCreateIndex(tx, "web_logs", "idx_web_logs_ai_score_day",
"CREATE INDEX IF NOT EXISTS idx_web_logs_ai_score_day ON web_logs (ai_score, day)")
},
Rollback: func(tx *gorm.DB) error {
zlog.Info("回滚 202606120002: 删除 web_logs.ai_score 索引")
return safeDropIndex(tx, "web_logs", "idx_web_logs_ai_score_day")
},
},
})
// 执行迁移
wafenginecore/checkai.go
+81
View File
@@ -0,0 +1,81 @@
package wafenginecore
import (
"SamWaf/global"
"SamWaf/innerbean"
"SamWaf/model/detection"
"SamWaf/model/wafenginmodel"
"net/http"
"net/url"
)
/*
*
AI智能检测
补现有规则引擎（正则 / libinjection / OWASP CRS）的盲区：变形绕过、混淆、未知 payload。
推理走纯 Go 内嵌 GBDT（wafai.Detector + dmitryikh/leaves），失败安全：
模型未加载/打分异常一律放行，绝不影响业务转发。
进入条件（在 wafengine 检测链中由调用方判断）：
- 全局开关 global.GCONFIG_AI_ENABLE == 1
- 站点开关 hostDefense.DEFENSE_AI == 1
工作模式 global.GCONFIG_AI_MODE：
- "observe"：仅记录命中（达到观察阈值即在日志标注 AI 分数），不拦截
- "block" ：达到模型建议拦截阈值则拦截，介于观察/拦截阈值之间则仅记录
*/
func (waf *WafEngine) CheckAI(r *http.Request, weblogbean *innerbean.WebLog, formValue url.Values, hostTarget *wafenginmodel.HostSafe, globalHostTarget *wafenginmodel.HostSafe) detection.Result {
result := detection.Result{
JumpGuardResult: false,
IsBlock: false,
Title: "",
Content: "",
}
detector := global.GWAF_AI_DETECTOR
if detector == nil || !detector.IsLoaded() {
return result
}
// 取请求关键字段做特征：path 与 query 分离，body 用解密后的明文
path := ""
if r != nil && r.URL != nil {
path = r.URL.Path
}
body := weblogbean.BODY
if body == "" {
body = weblogbean.POST_FORM
}
pred := detector.PredictRequest(weblogbean.METHOD, path, weblogbean.RawQuery, body, weblogbean.USER_AGENT)
if !pred.Loaded {
// 失败安全：未命中
return result
}
// 低于观察阈值：忽略（不记录分数，避免污染正常日志）
if pred.Score < pred.ObserveThreshold {
return result
}
// RULE 用稳定的类别标签（不含分数），便于"按规则汇总"统计；分数单独存 AI_SCORE 列
title := "AI检测:" + pred.Category
weblogbean.RISK_LEVEL = 2
weblogbean.AI_SCORE = pred.Score
// block 模式且达到拦截阈值 -> 拦截（由 handleBlock 走 EchoErrorInfo 记录为"阻止"）
if global.GCONFIG_AI_MODE == "block" && pred.Score >= pred.BlockThreshold {
result.IsBlock = true
result.Title = title
result.Content = "请正确访问"
return result
}
// 观察命中（observe 模式，或 block 模式下分数在[观察,拦截)之间）：
// 标记为"仅记录"，让请求照常放行，但在访问日志里可按 log_only_mode 筛出、
// 并通过 ai_score 列排序/查看——对标 AWS WAF Count / Cloudflare log 动作。
weblogbean.RULE = title
weblogbean.LogOnlyMode = 1
return result
}
wafenginecore/wafengine.go
+11
View File
@@ -89,6 +89,11 @@ func inferAttackType(ruleTitle string) string {
return "owasp_attack"
}
// AI 智能检测：Title 格式为 "AI检测:score=x.xx"，需在 SQL/RCE 等关键词匹配之前优先处理
if strings.HasPrefix(ruleTitle, "ai检测") {
return "ai_attack"
}
// CC攻击
if strings.Contains(ruleTitle, "cc") || strings.Contains(ruleTitle, "频次") || strings.Contains(ruleTitle, "rate limit") {
return "cc_attack"
@@ -556,6 +561,12 @@ func (waf *WafEngine) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if handleBlock(waf.CheckRule) {
return
}
//AI智能检测（全局开关开启且站点开启，规则抓确定的，AI 抓漏网的）
if global.GCONFIG_AI_ENABLE == 1 && hostDefense.DEFENSE_AI == 1 {
if handleBlock(waf.CheckAI) {
return
}
}
//检测敏感词
if hostDefense.DEFENSE_SENSITIVE == 1 {
if handleBlock(waf.CheckSensitive) {
wafmangeweb/localserver.go
+2
View File
@@ -173,6 +173,8 @@ func (web *WafWebManager) initRouter(r *gin.Engine) {
router.ApiGroupApp.InitSqlQueryRouter(TokenOnlyRouterGroup)
// 应用管理：可执行任意命令，拒绝 API Key，功能默认关闭
router.ApiGroupApp.InitWafAppRouter(TokenOnlyRouterGroup)
// AI模型管理与训练数据导出：模型会被引擎加载、数据出库，拒绝 API Key
router.ApiGroupApp.InitWafAIRouter(TokenOnlyRouterGroup)
}
// 保存 gin.Engine 引用供 API 文档生成使用
waftask/task_config.go
+14
View File
@@ -77,6 +77,9 @@ func setConfigIntValue(name string, value int64, change int) {
case "enable_owasp":
global.GCONFIG_RECORD_ENABLE_OWASP = value
break
case "ai_enable":
global.GCONFIG_AI_ENABLE = value
break
case "owasp_block_threshold":
if value <= 0 {
value = 7
@@ -228,6 +231,15 @@ func setConfigStringValue(name string, value string, change int) {
// 同步到 wafowasp 热路径，使 DetectionOnly "本该拦截" 的 INFO 日志能按当前模式生效
wafowasp.SetEngineMode(global.GCONFIG_OWASP_MODE)
break
case "ai_mode":
switch value {
case "observe", "block":
global.GCONFIG_AI_MODE = value
default:
zlog.Warn("invalid ai_mode value, fallback to observe", value)
global.GCONFIG_AI_MODE = "observe"
}
break
case "kafka_url":
global.GCONFIG_RECORD_KAFKA_URL = value
break
@@ -394,6 +406,8 @@ func TaskLoadSetting(initLoad bool) {
updateConfigIntItem(initLoad, "system", "login_max_error_time", global.GCONFIG_RECORD_LOGIN_MAX_ERROR_TIME, "登录周期里错误最大次数 请大于0 ", "int", "", configMap)
updateConfigIntItem(initLoad, "system", "login_limit_mintutes", global.GCONFIG_RECORD_LOGIN_LIMIT_MINTUTES, "登录错误记录周期 单位分钟数，默认1分钟", "int", "", configMap)
updateConfigIntItem(initLoad, "system", "enable_owasp", global.GCONFIG_RECORD_ENABLE_OWASP, "启动OWASP数据检测（1启动 0关闭）", "int", "", configMap)
updateConfigIntItem(initLoad, "system", "ai_enable", global.GCONFIG_AI_ENABLE, "启动AI智能检测总开关（1启动 0关闭，需先在AI模型管理上传模型包并在站点开启）", "int", "", configMap)
updateConfigStringItem(initLoad, "system", "ai_mode", global.GCONFIG_AI_MODE, "AI检测工作模式：observe(仅记录/观察) block(达拦截阈值则拦截)", "options", "observe|仅记录,block|拦截", configMap)
updateConfigIntItem(initLoad, "ssl", "enable_http_80", global.GCONFIG_RECORD_ENABLE_HTTP_80, "启动80端口服务（为自动申请证书使用 HTTP文件验证类型需要，DNS验证不需要）", "int", "", configMap)
updateConfigIntItem(initLoad, "ssl", "sslorder_expire_day", global.GCONFIG_RECORD_SSLOrder_EXPIRE_DAY, "自动续期检测小于多少天开始发起自动申请 默认30天", "int", "", configMap)