diff --git a/api/entrance.go b/api/entrance.go index 9bf5d0f..7e856a1 100644 --- a/api/entrance.go +++ b/api/entrance.go @@ -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 diff --git a/api/waf_ai_api.go b/api/waf_ai_api.go new file mode 100644 index 0000000..153aa41 --- /dev/null +++ b/api/waf_ai_api.go @@ -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) +} diff --git a/cmd/samwaf/main.go b/cmd/samwaf/main.go index 8d7c214..c07b2f4 100644 --- a/cmd/samwaf/main.go +++ b/cmd/samwaf/main.go @@ -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 { diff --git a/global/config.go b/global/config.go index 10967b7..13c575a 100644 --- a/global/config.go +++ b/global/config.go @@ -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 //记录原始信息(默认不开启) diff --git a/global/global.go b/global/global.go index 78c5ef5..617bd0f 100644 --- a/global/global.go +++ b/global/global.go @@ -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 管理器（支持热重载） diff --git a/go.mod b/go.mod index 9d63ffa..9c3316f 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 3e54055..b0c0d17 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/innerbean/web_log.go b/innerbean/web_log.go index 00f5d13..9386ff8 100644 --- a/innerbean/web_log.go +++ b/innerbean/web_log.go @@ -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的值 diff --git a/model/hosts.go b/model/hosts.go index 43e7ff3..065b77c 100644 --- a/model/hosts.go +++ b/model/hosts.go @@ -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 != "" { diff --git a/model/request/waf_ai_req.go b/model/request/waf_ai_req.go new file mode 100644 index 0000000..b3bd983 --- /dev/null +++ b/model/request/waf_ai_req.go @@ -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"` +} diff --git a/model/response/waf_ai_resp.go b/model/response/waf_ai_resp.go new file mode 100644 index 0000000..577f6ba --- /dev/null +++ b/model/response/waf_ai_resp.go @@ -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"` +} diff --git a/model/waf_log_label_mark.go b/model/waf_log_label_mark.go new file mode 100644 index 0000000..f0f28f6 --- /dev/null +++ b/model/waf_log_label_mark.go @@ -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（快照） +} diff --git a/router/entrance.go b/router/entrance.go index 0afc709..be97a3b 100644 --- a/router/entrance.go +++ b/router/entrance.go @@ -59,6 +59,7 @@ type ApiGroup struct { WafOwaspRouter WafHostPathRuleRouter WafAppRouter + WafAIRouter } type PublicApiGroup struct { LoginRouter diff --git a/router/waf_ai_router.go b/router/waf_ai_router.go new file mode 100644 index 0000000..3208c49 --- /dev/null +++ b/router/waf_ai_router.go @@ -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) + } +} diff --git a/service/waf_service/waf_ai_label_service.go b/service/waf_service/waf_ai_label_service.go new file mode 100644 index 0000000..99a3a7b --- /dev/null +++ b/service/waf_service/waf_ai_label_service.go @@ -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 +} diff --git a/service/waf_service/waf_ai_service.go b/service/waf_service/waf_ai_service.go new file mode 100644 index 0000000..a80a536 --- /dev/null +++ b/service/waf_service/waf_ai_service.go @@ -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 +} diff --git a/wafai/engine.go b/wafai/engine.go new file mode 100644 index 0000000..0d55361 --- /dev/null +++ b/wafai/engine.go @@ -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 +} diff --git a/wafai/engine_gbdt.go b/wafai/engine_gbdt.go new file mode 100644 index 0000000..d44f6d8 --- /dev/null +++ b/wafai/engine_gbdt.go @@ -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" } diff --git a/wafai/engine_onnx_stub.go b/wafai/engine_onnx_stub.go new file mode 100644 index 0000000..32342a7 --- /dev/null +++ b/wafai/engine_onnx_stub.go @@ -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 编译）") +} diff --git a/wafai/features.go b/wafai/features.go new file mode 100644 index 0000000..3964f44 --- /dev/null +++ b/wafai/features.go @@ -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(" 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) +} diff --git a/wafai/features_test.go b/wafai/features_test.go new file mode 100644 index 0000000..6791457 --- /dev/null +++ b/wafai/features_test.go @@ -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=", "", "", "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"))) + } +} diff --git a/wafai/labeling.go b/wafai/labeling.go new file mode 100644 index 0000000..5c82dd4 --- /dev/null +++ b/wafai/labeling.go @@ -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, "" +} diff --git a/wafai/manager.go b/wafai/manager.go new file mode 100644 index 0000000..aaaf1c8 --- /dev/null +++ b/wafai/manager.go @@ -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, + } +} diff --git a/wafai/modelpack.go b/wafai/modelpack.go new file mode 100644 index 0000000..6d757b4 --- /dev/null +++ b/wafai/modelpack.go @@ -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 +} diff --git a/wafai/modelpack_test.go b/wafai/modelpack_test.go new file mode 100644 index 0000000..9ab0c1b --- /dev/null +++ b/wafai/modelpack_test.go @@ -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("含路径穿越条目应加载失败") + } +} diff --git a/wafai/testdata/feature_golden.json b/wafai/testdata/feature_golden.json new file mode 100644 index 0000000..6f2f3c8 --- /dev/null +++ b/wafai/testdata/feature_golden.json @@ -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 + ] + } + ] +} \ No newline at end of file diff --git a/wafai/testdata/sample_model.swai b/wafai/testdata/sample_model.swai new file mode 100644 index 0000000..04077bb Binary files /dev/null and b/wafai/testdata/sample_model.swai differ diff --git a/wafai/testhelper_test.go b/wafai/testhelper_test.go new file mode 100644 index 0000000..d8fa11d --- /dev/null +++ b/wafai/testhelper_test.go @@ -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() +} diff --git a/wafdb/migrations_core.go b/wafdb/migrations_core.go index 8f53e11..188c14c 100644 --- a/wafdb/migrations_core.go +++ b/wafdb/migrations_core.go @@ -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 + }, + }, }) // 执行迁移 diff --git a/wafdb/migrations_log.go b/wafdb/migrations_log.go index dd75ff2..ab729c3 100644 --- a/wafdb/migrations_log.go +++ b/wafdb/migrations_log.go @@ -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") + }, + }, }) // 执行迁移 diff --git a/wafenginecore/checkai.go b/wafenginecore/checkai.go new file mode 100644 index 0000000..642de2a --- /dev/null +++ b/wafenginecore/checkai.go @@ -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 +} diff --git a/wafenginecore/wafengine.go b/wafenginecore/wafengine.go index 19e3771..d7cccd2 100644 --- a/wafenginecore/wafengine.go +++ b/wafenginecore/wafengine.go @@ -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) { diff --git a/wafmangeweb/localserver.go b/wafmangeweb/localserver.go index 310514e..6af83dc 100644 --- a/wafmangeweb/localserver.go +++ b/wafmangeweb/localserver.go @@ -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 文档生成使用 diff --git a/waftask/task_config.go b/waftask/task_config.go index 21e7934..3da7b2d 100644 --- a/waftask/task_config.go +++ b/waftask/task_config.go @@ -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)