fix: 修复 POC Cookie/变量提取的两个问题

- poc_adapter: xray/afrog 的 output.search 转换为 Search 字段,多步POC变量传递不再丢失
- poc_executor: Set-Cookie 提取优化不再要求捕获组名含 cookie,sessid/token等命名均生效
This commit is contained in:
ZacharyZcR
2026-05-13 18:37:33 +08:00
parent 3436d6ad02
commit b2e91d9fc0
4 changed files with 258 additions and 2 deletions
+10
View File
@@ -399,6 +399,11 @@ func (x *XrayPocAdapter) ToFscanPoc() (*Poc, error) {
Expression: rule.Expression,
}
// 转换 output 字段为 Search — 多步POC中从响应提取变量供后续步骤使用
if searchVal, ok := rule.Output["search"]; ok {
fscanRule.Search = fmt.Sprintf("%v", searchVal)
}
// 如果expression为空,默认检查200状态码
if fscanRule.Expression == "" {
fscanRule.Expression = "response.status == 200"
@@ -500,6 +505,11 @@ func (a *AfrogPocAdapter) ToFscanPoc() (*Poc, error) {
Expression: rule.Expression,
}
// 转换 output 字段为 Search — 多步POC中从响应提取变量供后续步骤使用
if searchVal, ok := rule.Output["search"]; ok {
fscanRule.Search = fmt.Sprintf("%v", searchVal)
}
// 如果expression为空,默认检查200状态码
if fscanRule.Expression == "" {
fscanRule.Expression = "response.status == 200"
+149
View File
@@ -1,6 +1,7 @@
package lib
import (
"strings"
"testing"
)
@@ -383,3 +384,151 @@ unknown: format
})
}
}
// TestXrayOutputToSearch 测试 xray output 字段到 Search 的转换
func TestXrayOutputToSearch(t *testing.T) {
// 多步POC: r0 提取 cookier1 使用 {{cookie}}
yamlData := `
name: poc-yaml-test-cookie-extract
transport: http
rules:
r0:
request:
method: POST
path: /login
headers:
Content-Type: text/xml
body: userID=admin
follow_redirects: false
expression: response.status == 200
output:
search: "Set-Cookie:(?P<cookie>.*)"
r1:
request:
method: GET
path: /admin/dashboard
headers:
Cookie: "{{cookie}}"
expression: response.status == 200
detail:
author: test
`
adapter, err := loadXrayPoc([]byte(yamlData))
if err != nil {
t.Fatalf("loadXrayPoc() error = %v", err)
}
poc, err := adapter.ToFscanPoc()
if err != nil {
t.Fatalf("ToFscanPoc() error = %v", err)
}
if len(poc.Rules) != 2 {
t.Fatalf("len(Poc.Rules) = %d, want 2", len(poc.Rules))
}
// r0 应该有 Search 字段(从 output.search 转换)
if poc.Rules[0].Search == "" {
t.Error("Rules[0].Search should not be empty — output.search was not converted")
}
if !strings.Contains(poc.Rules[0].Search, "cookie") {
t.Errorf("Rules[0].Search = %q, should contain 'cookie'", poc.Rules[0].Search)
}
// r1 不应该有 Search(没有 output 字段)
if poc.Rules[1].Search != "" {
t.Errorf("Rules[1].Search = %q, should be empty", poc.Rules[1].Search)
}
// r1 的 Headers 应保留 {{cookie}} 占位符
if poc.Rules[1].Headers["Cookie"] != `{{cookie}}` {
t.Errorf("Rules[1].Headers[Cookie] = %q, want %q", poc.Rules[1].Headers["Cookie"], `{{cookie}}`)
}
}
// TestXrayNoOutput 测试 xray 没有 output 字段时 Search 为空(回归)
func TestXrayNoOutput(t *testing.T) {
yamlData := `
name: poc-yaml-test-simple
transport: http
rules:
r0:
request:
method: GET
path: /api/test
expression: response.status == 200
detail:
author: test
`
adapter, err := loadXrayPoc([]byte(yamlData))
if err != nil {
t.Fatalf("loadXrayPoc() error = %v", err)
}
poc, err := adapter.ToFscanPoc()
if err != nil {
t.Fatalf("ToFscanPoc() error = %v", err)
}
if len(poc.Rules) != 1 {
t.Fatalf("len(Poc.Rules) = %d, want 1", len(poc.Rules))
}
if poc.Rules[0].Search != "" {
t.Errorf("Rules[0].Search = %q, should be empty when no output field", poc.Rules[0].Search)
}
}
// TestAfrogOutputToSearch 测试 afrog output 字段到 Search 的转换
func TestAfrogOutputToSearch(t *testing.T) {
yamlData := `
id: test-afrog-cookie
info:
name: 测试Cookie提取
author: test
severity: high
rules:
r0:
request:
method: POST
path: /login
headers:
Content-Type: application/x-www-form-urlencoded
body: username=admin&password=123456
expression: response.status == 200 && response.body.bcontains(b"success")
output:
search: "Set-Cookie:(?P<sessid>.*)"
r1:
request:
method: GET
path: /panel
headers:
Cookie: "{{sessid}}"
expression: response.status == 200 && response.body.bcontains(b"admin")
`
adapter, err := loadAfrogPoc([]byte(yamlData))
if err != nil {
t.Fatalf("loadAfrogPoc() error = %v", err)
}
poc, err := adapter.ToFscanPoc()
if err != nil {
t.Fatalf("ToFscanPoc() error = %v", err)
}
if len(poc.Rules) != 2 {
t.Fatalf("len(Poc.Rules) = %d, want 2", len(poc.Rules))
}
if poc.Rules[0].Search == "" {
t.Error("Rules[0].Search should not be empty — output.search was not converted")
}
// r1 的占位符应对应捕获组名 sessid
if poc.Rules[1].Headers["Cookie"] != `{{sessid}}` {
t.Errorf("Rules[1].Headers[Cookie] = %q, want %q", poc.Rules[1].Headers["Cookie"], `{{sessid}}`)
}
}
+2 -2
View File
@@ -348,8 +348,8 @@ func doSearch(re string, body string) map[string]string {
paramsMap := make(map[string]string)
for i, name := range names {
if i > 0 && i <= len(result) {
// 特殊处理Cookie头
if strings.HasPrefix(re, "Set-Cookie:") && strings.Contains(name, "cookie") {
// 特殊处理Set-Cookie头:剥离Path/Expires等属性,仅保留key=value
if strings.HasPrefix(re, "Set-Cookie:") {
paramsMap[name] = optimizeCookies(result[i])
} else {
paramsMap[name] = result[i]
+97
View File
@@ -1,6 +1,7 @@
package lib
import (
"strings"
"testing"
)
@@ -101,6 +102,102 @@ func TestGetRuleHash(t *testing.T) {
}
}
// TestDoSearchSetCookieOptimization 测试 Set-Cookie 提取和清理
func TestDoSearchSetCookieOptimization(t *testing.T) {
responseHeaders := "HTTP/1.1 200 OK\r\n"
cases := []struct {
name string
regex string
body string
wantContain string // 期望结果包含的内容
wantNotContain string // 期望结果不包含的内容
}{
{
name: "捕获组名为cookie时清理属性",
regex: `Set-Cookie:(?P<cookie>.*)`,
body: responseHeaders + "Set-Cookie: sessionid=abc123; Path=/; HttpOnly\r\n\r\n<html></html>",
wantContain: "sessionid=abc123",
wantNotContain: "Path",
},
{
name: "捕获组名为sessid时也清理属性",
regex: `Set-Cookie:(?P<sessid>.*)`,
body: responseHeaders + "Set-Cookie: JSESSIONID=xyz789; Path=/app; Secure; HttpOnly\r\n\r\n{}",
wantContain: "JSESSIONID=xyz789",
wantNotContain: "Secure",
},
{
name: "捕获组名为token时也清理属性",
regex: `Set-Cookie:(?P<token>.*)`,
body: responseHeaders + "Set-Cookie: csrf_token=tok123; Max-Age=3600; SameSite=Strict\r\n\r\nOK",
wantContain: "csrf_token=tok123",
wantNotContain: "Max-Age",
},
{
name: "非Set-Cookie的正则不触发清理",
regex: `X-Custom:(?P<value>.*)`,
body: responseHeaders + "X-Custom: some-value; extra=stuff\r\n\r\ndone",
wantContain: "some-value; extra=stuff",
},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
result := doSearch(c.regex, c.body)
if result == nil {
t.Fatal("doSearch() returned nil")
}
for _, v := range result {
if c.wantContain != "" && !strings.Contains(v, c.wantContain) {
t.Errorf("result should contain %q, got %q", c.wantContain, v)
}
if c.wantNotContain != "" && strings.Contains(v, c.wantNotContain) {
t.Errorf("result should NOT contain %q, got %q", c.wantNotContain, v)
}
}
})
}
}
// TestOptimizeCookies 测试 Cookie 清理函数
func TestOptimizeCookies(t *testing.T) {
cases := []struct {
name string
raw string
want string
}{
{
name: "标准Set-Cookie带多个属性",
raw: "sessionid=abc123; Path=/; HttpOnly; Secure",
want: "sessionid=abc123",
},
{
name: "多个cookie键值对",
raw: "token=xyz; user=admin; Path=/app; Expires=Wed, 21 Oct 2025 07:28:00 GMT",
want: "token=xyz; user=admin",
},
{
name: "无属性的干净cookie",
raw: "sid=simple",
want: "sid=simple",
},
{
name: "空字符串",
raw: "",
want: "",
},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
got := optimizeCookies(c.raw)
if got != c.want {
t.Errorf("optimizeCookies(%q) = %q, want %q", c.raw, got, c.want)
}
})
}
}
// TestApplyParametersToRule 测试参数替换逻辑
func TestApplyParametersToRule(t *testing.T) {
tests := []struct {