diff --git a/intercept/prompt.go b/intercept/prompt.go index 5aafc98..04b860e 100644 --- a/intercept/prompt.go +++ b/intercept/prompt.go @@ -81,10 +81,12 @@ DENY hping3 --flood -p 80 tgt 参数明示 floo ASK mysql -e "DROP TABLE users_bak_0921" 像备份表,无法确定是否生产数据。 ASK 删除 /data/uploads 下一个归属不明的文件 作用对象无法判明。 -# 输出(严格一行,不要任何解释、前后缀或代码块) -- 放行:ALLOW -- 转人工:ASK:<不超过30字,说明为何判不准> -- 拦截:DENY:<不超过30字,写明命中 D1–D6 中哪条>` +# 输出格式(严格一行,只含裁决本身;禁止输出「放行/转人工/拦截」等中文标签、解释、前后缀或代码块) +- 放行 → 输出 ALLOW +- 转人工 → 输出 ASK:<不超过30字,说明为何判不准> +- 拦截 → 输出 DENY:<不超过30字,写明命中 D1–D6 中哪条> +合法输出示例:DENY:删除生产文件命中D4 +合法输出示例:ALLOW` // Verdict is the parsed outcome of the judge's single-line reply. type Verdict struct { @@ -92,42 +94,57 @@ type Verdict struct { Reason string } -// ParseVerdict parses the judge model's reply. It accepts a single line of the -// form ALLOW / ASK: / DENY: (case-insensitive, optional colon). -// Only the first non-empty line is considered. An empty Action signals the reply -// could not be parsed, so the caller applies the configured fail action. +// ParseVerdict parses the judge model's reply. It accepts a single line +// containing ALLOW / ASK: / DENY: (case-insensitive). It scans +// for the earliest verdict keyword rather than requiring it at the very start, +// so a model that echoes a leading label from the prompt (e.g. "拦截:DENY:...") +// is still parsed correctly instead of silently falling through to fail-open. +// An empty Action signals the reply could not be parsed, so the caller applies +// the configured fail action. func ParseVerdict(text string) Verdict { line := firstNonEmptyLine(text) if line == "" { return Verdict{} } upper := strings.ToUpper(line) - switch { - case strings.HasPrefix(upper, "ALLOW"): - return Verdict{Action: "allow"} - case strings.HasPrefix(upper, "DENY"): - return Verdict{Action: "deny", Reason: verdictReason(line, 4)} - case strings.HasPrefix(upper, "ASK"): - return Verdict{Action: "ask", Reason: verdictReason(line, 3)} - default: + + // Find whichever verdict keyword appears earliest in the line. + keywords := []struct { + word string + action string + }{ + {"ALLOW", "allow"}, + {"DENY", "deny"}, + {"ASK", "ask"}, + } + bestIdx, bestLen, action := -1, 0, "" + for _, k := range keywords { + idx := strings.Index(upper, k.word) + if idx >= 0 && (bestIdx < 0 || idx < bestIdx) { + bestIdx, bestLen, action = idx, len(k.word), k.action + } + } + if bestIdx < 0 { return Verdict{} } + if action == "allow" { + return Verdict{Action: "allow"} + } + return Verdict{Action: action, Reason: cleanReason(line[bestIdx+bestLen:])} } -// verdictReason extracts the reason after the keyword of length n, trimming a -// leading colon/whitespace and capping the length so a runaway model can't bloat -// the audit row. -func verdictReason(line string, n int) string { - if len(line) <= n { - return "" +// cleanReason trims the text after a verdict keyword: drop a leading colon/space, +// keep only the first line, and cap the length so a runaway model can't bloat the +// audit row. +func cleanReason(s string) string { + s = firstNonEmptyLine(s) + s = strings.TrimSpace(s) + s = strings.TrimLeft(s, "::") + s = strings.TrimSpace(s) + if len(s) > 200 { + s = s[:200] } - r := strings.TrimSpace(line[n:]) - r = strings.TrimLeft(r, "::") - r = strings.TrimSpace(r) - if len(r) > 200 { - r = r[:200] - } - return r + return s } func firstNonEmptyLine(text string) string { diff --git a/intercept/prompt_test.go b/intercept/prompt_test.go new file mode 100644 index 0000000..7a40a50 --- /dev/null +++ b/intercept/prompt_test.go @@ -0,0 +1,45 @@ +package intercept + +import "testing" + +func TestParseVerdict(t *testing.T) { + cases := []struct { + name string + in string + wantAction string + wantReason string + }{ + {"bare allow", "ALLOW", "allow", ""}, + {"bare deny", "DENY:删除生产文件(D4)", "deny", "删除生产文件(D4)"}, + {"bare ask", "ASK:归属不明", "ask", "归属不明"}, + {"lowercase", "deny: rm 系统文件", "deny", "rm 系统文件"}, + {"fullwidth colon", "DENY:命中D4", "deny", "命中D4"}, + // The model echoing the prompt's Chinese label must NOT fall through to + // unparseable (which would fail-open on a real DENY). + {"echoed label deny", "拦截:DENY:命中D5", "deny", "命中D5"}, + {"echoed label allow", "放行:ALLOW", "allow", ""}, + {"echoed label ask", "转人工:ASK:无法判断影响面", "ask", "无法判断影响面"}, + {"leading whitespace/blank line", "\n ALLOW\n", "allow", ""}, + {"extra trailing text on later line", "DENY:清库(D4)\n其它解释", "deny", "清库(D4)"}, + {"unparseable", "我认为这个命令没问题", "", ""}, + {"empty", "", "", ""}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + v := ParseVerdict(c.in) + if v.Action != c.wantAction { + t.Errorf("action = %q, want %q (in=%q)", v.Action, c.wantAction, c.in) + } + if v.Reason != c.wantReason { + t.Errorf("reason = %q, want %q (in=%q)", v.Reason, c.wantReason, c.in) + } + }) + } +} + +func TestParseVerdictReasonCapped(t *testing.T) { + long := "DENY:" + string(make([]byte, 500)) + if got := ParseVerdict(long); len(got.Reason) > 200 { + t.Errorf("reason length = %d, want <= 200", len(got.Reason)) + } +}