mirror of
https://github.com/RuoJi6/ARTEX.git
synced 2026-10-27 17:03:49 +08:00
fix(agent): list_assets 按任务范围过滤,修复任务间资产互串 (#59)
list_assets 此前把任务 id 写死为 0、裸查整个共享资产库,模型又无从写出 范围过滤条件,于是把别的任务的资产(尤其 IP)当成本任务目标,测试方向 被带偏。 - 新增 QueryDSLInScope / GetByIDsInScope:只返回落在【本任务及直接关联 任务】task_scope 内的资产,按归属匹配而非字面值(根域范围含其全部子域/ 服务/接口,网段范围含段内主机及服务);越权 id 静默丢弃 - 新增 scopeTargetCTE(纯 scope 归属判定,不含 anchor/tested);whereBuilder 加占位符偏移,$1 留给 taskID - 修复 IP 直连主机(http://1.2.3.4/api)无法被网段范围命中的盲区:谓词加 OR ts.net >>= try_inet(a.domain) - 非任务上下文(Auto/pentest, taskID<=0)无范围可依,回退全库 - QueryDSL/CountDSL 原样未动,UI 的「测试资产」视图(按生产者过滤)不受影响
This commit is contained in:
1 parent
16e87b48c2
commit
4ccb8f4ff2
5 files changed
+267
-7
No files matched your search
@@ -11,6 +11,10 @@
|
||||
- `insert_assets` 移除每条资产的 `related` 入参:它唯一的作用是决定要不要把这条资产加进本任务范围,值又不落库(同一资产被再次登记就作废、界面上也看不出谁被判为无关),实际只是让模型多做一次留不住的判断。
|
||||
- **测试范围(`task_scope`)不再受资产覆盖度开关影响**:`insert_assets` 的自动入范围(`source='auto'`)无论覆盖度开启与否都执行,`add_task_scope` 也始终提供给 planner / 任务内主 Agent / 目标拆解 Agent。范围是任务的授权边界、资产查询的过滤基准,覆盖度开关只决定要不要拿它当分母算指标,不该决定要不要累积范围本身。此前关闭覆盖度会让 `auto`/`agent` 两条写入路径同时失效,`task_scope` 只剩界面手工添加的行。`list_untested_assets` 仍随覆盖度关闭而隐藏(它本身就是纯覆盖度视角)。
|
||||
|
||||
#### 修复的问题
|
||||
|
||||
- **修复任务间资产互串**(#59):`list_assets` 此前把任务 id 写死为 0、裸查整个共享资产库,模型又无从写出范围过滤条件,于是把别的任务的资产(尤其 IP)当成本任务目标,测试方向被带偏。现在它只返回落在【本任务及直接关联任务】测试范围(`task_scope`)内的资产,按**归属**匹配而非字面值:范围里有某根域即可查到其名下全部子域/服务/接口,有某网段即可查到段内主机及服务;按 id 直取范围外资产同样取不到。非任务上下文(Auto/pentest)无范围可依,仍回退全库。顺带修复 IP 直连主机(如 `http://1.2.3.4/api`)无法被网段范围命中的归属盲区。UI 的「测试资产」视图(按任务生产者过滤)不受影响。
|
||||
|
||||
## [0.3.9] - 2026-09-11
|
||||
|
||||
### Agent
|
||||
|
||||
@@ -456,10 +456,10 @@ func (t *ToolSet) listUntestedAssets() actool.CoreTool {
|
||||
func (t *ToolSet) listAssets() actool.CoreTool {
|
||||
return readTool(
|
||||
"list_assets",
|
||||
"查询资产库:DSL 表达式搜索,或按 id/ids 直取;支持分页。\n"+
|
||||
"查询资产库:DSL 表达式搜索,或按 id/ids 直取;支持分页。只返回【本任务及直接关联任务】测试范围内的资产。\n"+
|
||||
"DSL:field=value 模糊(ILIKE) | field==value 精确 | field!=value 排除 | 数字字段支持 > >= < <= | 裸词=全文模糊;AND/OR 组合(AND 优先级高),可用括号分组。资产类型用独立 type 参数,不写进 DSL。\n"+
|
||||
"未传 id/ids 时 dsl 必须非空(不允许无条件全量查询)。\n"+
|
||||
"可用字段:domain(根/子/服务域名)、root_domain、ip、url、page_title、icp、service_name、app_name、method(如 GET/POST)、service_type(http|other)、record_type(如 A/CNAME)、technology(数组,=模糊 ==精确)、port/status_code/company_id/task_id(整数)。\n"+
|
||||
"可用字段:domain(根/子/服务域名)、root_domain、ip、url、page_title、icp、service_name、app_name、method(如 GET/POST)、service_type(http|other)、record_type(如 A/CNAME)、technology(数组,=模糊 ==精确)、port/status_code/company_id(整数)。\n"+
|
||||
"示例:status_code>=400 AND technology=shiro ;(port==80 OR port==443) AND technology=nginx",
|
||||
obj(map[string]any{
|
||||
"dsl": str(`DSL 查询表达式(语法/字段见工具描述)。未传 id/ids 时必须非空。`),
|
||||
@@ -490,11 +490,11 @@ func (t *ToolSet) listAssets() actool.CoreTool {
|
||||
var err error
|
||||
switch {
|
||||
case a.ID > 0:
|
||||
assets, err = t.as.GetByIDs([]int64{a.ID})
|
||||
assets, err = t.as.GetByIDsInScope(t.taskID, []int64{a.ID})
|
||||
case len(a.IDs) > 0:
|
||||
assets, err = t.as.GetByIDs(a.IDs)
|
||||
assets, err = t.as.GetByIDsInScope(t.taskID, a.IDs)
|
||||
case a.DSL != "":
|
||||
assets, err = t.as.QueryDSL(a.DSL, a.Type, 0, a.Limit, a.Offset)
|
||||
assets, err = t.as.QueryDSLInScope(a.DSL, a.Type, t.taskID, a.Limit, a.Offset)
|
||||
default:
|
||||
return actool.Errorf("未传 id/ids 时 dsl 不能为空:不允许无条件查询全部资产,请提供查询条件"), nil
|
||||
}
|
||||
|
||||
+97
-2
@@ -259,11 +259,12 @@ var fullTextCols = []string{
|
||||
|
||||
type whereBuilder struct {
|
||||
args []any
|
||||
base int // placeholders are numbered base+1, base+2, …; 0 = the usual $1, $2, …
|
||||
}
|
||||
|
||||
func (b *whereBuilder) next(v any) string {
|
||||
b.args = append(b.args, v)
|
||||
return fmt.Sprintf("$%d", len(b.args))
|
||||
return fmt.Sprintf("$%d", b.base+len(b.args))
|
||||
}
|
||||
|
||||
func (b *whereBuilder) build(node *astNode) (string, error) {
|
||||
@@ -378,10 +379,17 @@ func (b *whereBuilder) buildLeaf(e Expr) (string, error) {
|
||||
}
|
||||
|
||||
func buildDSLWhere(node *astNode) (string, []any, error) {
|
||||
return buildDSLWhereBase(node, 0)
|
||||
}
|
||||
|
||||
// buildDSLWhereBase is buildDSLWhere with a placeholder offset: emitted args are
|
||||
// numbered base+1 onward, leaving $1..$base free for the caller (e.g. a scope CTE
|
||||
// that reserves $1 for the task id).
|
||||
func buildDSLWhereBase(node *astNode, base int) (string, []any, error) {
|
||||
if node == nil {
|
||||
return "1=1", nil, nil
|
||||
}
|
||||
b := &whereBuilder{}
|
||||
b := &whereBuilder{base: base}
|
||||
clause, err := b.build(node)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
@@ -540,3 +548,90 @@ func (s *AssetStore) QueryDSL(dsl, typ string, taskID int64, limit, offset int)
|
||||
}
|
||||
return assets, nil
|
||||
}
|
||||
|
||||
// QueryDSLInScope is QueryDSL restricted to assets that BELONG to taskID's (and its
|
||||
// direct source tasks') declared scope — membership, not literal value: a
|
||||
// root_domain scope returns every subdomain / service / endpoint under it. This is
|
||||
// the agent-facing list_assets path, so an agent queries the task's relevant assets
|
||||
// instead of the whole shared库. taskID<=0 (non-task contexts: Auto / pentest / chat)
|
||||
// has no scope to honor and falls back to the plain global QueryDSL. Rows carry the
|
||||
// same per-task source metadata as QueryByTask.
|
||||
func (s *AssetStore) QueryDSLInScope(dsl, typ string, taskID int64, limit, offset int) ([]*Asset, error) {
|
||||
if taskID <= 0 {
|
||||
return s.QueryDSL(dsl, typ, 0, limit, offset)
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
node, err := ParseDSL(dsl)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// $1 is reserved for taskID (scopeTargetCTE); DSL placeholders start at $2.
|
||||
where, dslArgs, err := buildDSLWhereBase(node, 1)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
args := []any{taskID}
|
||||
args = append(args, dslArgs...)
|
||||
if typ != "" {
|
||||
args = append(args, typ)
|
||||
where += fmt.Sprintf(" AND type = $%d", len(args))
|
||||
}
|
||||
where += " AND id IN (SELECT id FROM target)"
|
||||
args = append(args, limit, offset)
|
||||
q := `WITH ` + scopeTargetCTE + ` ` + assetSelectCols + " WHERE " + where +
|
||||
fmt.Sprintf(" ORDER BY last_seen DESC, id DESC LIMIT $%d OFFSET $%d", len(args)-1, len(args))
|
||||
rows, err := s.db.Query(q, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
assets, err := scanAssets(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.hydrateTaskAssetSources(taskID, assets); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return assets, nil
|
||||
}
|
||||
|
||||
// GetByIDsInScope is GetByIDs restricted to ids that BELONG to taskID's (and its
|
||||
// direct source tasks') declared scope, so an agent cannot reach out-of-scope
|
||||
// assets by id. taskID<=0 (non-task contexts) falls back to the global GetByIDs.
|
||||
// Out-of-scope ids are silently dropped from the result (not an error).
|
||||
func (s *AssetStore) GetByIDsInScope(taskID int64, ids []int64) ([]*Asset, error) {
|
||||
if taskID <= 0 {
|
||||
return s.GetByIDs(ids)
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
args := make([]any, 0, len(ids)+1)
|
||||
args = append(args, taskID) // $1 reserved for scopeTargetCTE
|
||||
placeholders := make([]string, len(ids))
|
||||
for i, id := range ids {
|
||||
placeholders[i] = fmt.Sprintf("$%d", i+2)
|
||||
args = append(args, id)
|
||||
}
|
||||
q := `WITH ` + scopeTargetCTE + ` ` + assetSelectCols +
|
||||
" WHERE id IN (" + strings.Join(placeholders, ",") + ")" +
|
||||
" AND id IN (SELECT id FROM target) ORDER BY last_seen DESC, id DESC"
|
||||
rows, err := s.db.Query(q, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
assets, err := scanAssets(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.hydrateTaskAssetSources(taskID, assets); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return assets, nil
|
||||
}
|
||||
@@ -999,3 +999,128 @@ func TestAssetIPRejectsHostname(t *testing.T) {
|
||||
}
|
||||
deleteAsset(d, id)
|
||||
}
|
||||
|
||||
// TestQueryDSLInScopeMembership pins the agent-facing list_assets behavior: it
|
||||
// returns assets that BELONG to the task's declared scope (membership, not literal
|
||||
// value) — a root_domain scope surfaces its subdomains/services/endpoints even when
|
||||
// those rows were produced by a different task — and honors ip/cidr scope for
|
||||
// IP-literal hosts (the try_inet(domain) arm). Out-of-scope assets, including by id,
|
||||
// are excluded. Direct source tasks' scope is included.
|
||||
func TestQueryDSLInScopeMembership(t *testing.T) {
|
||||
d, assets, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
|
||||
task, err := d.CreateTask("scope membership", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
src, err := d.CreateTask("scope source", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID); _ = d.DeleteTask(src.ID) })
|
||||
// task ← src is a direct source relation.
|
||||
if _, err := d.Exec(`INSERT INTO task_relations(task_id, source_task_id) VALUES ($1,$2)`, task.ID, src.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
stamp := time.Now().UnixNano()
|
||||
root := fmt.Sprintf("sc%d.invalid", stamp) // in-scope root domain (task)
|
||||
srcRoot := fmt.Sprintf("src%d.invalid", stamp) // in-scope via source task
|
||||
out := fmt.Sprintf("out%d.invalid", stamp) // out of every scope
|
||||
marker := fmt.Sprintf("mk%d", stamp) // bare-text token present in all rows
|
||||
foreignTask := stamp + 777 // asset produced by an unrelated task
|
||||
|
||||
// Scope: task owns root; source task owns srcRoot; task owns an IP /24.
|
||||
for _, sc := range []struct {
|
||||
tid int64
|
||||
kind string
|
||||
dom string
|
||||
net string
|
||||
}{
|
||||
{task.ID, "root_domain", root, ""},
|
||||
{src.ID, "root_domain", srcRoot, ""},
|
||||
{task.ID, "ip", "", "198.51.100.0/24"},
|
||||
} {
|
||||
if err := assets.upsertTaskScope(TaskScope{TaskID: sc.tid, Kind: sc.kind, Domain: sc.dom, Net: sc.net}); err != nil {
|
||||
t.Fatalf("seed scope: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Insert rows directly so root_domain/task_ids are controlled exactly. Every
|
||||
// in-scope row carries foreignTask in task_ids (never the current task) to prove
|
||||
// membership is by scope, not by producer.
|
||||
type row struct {
|
||||
typ, domain, rootDom, url, method, ip, title string
|
||||
}
|
||||
rows := []row{
|
||||
{"subdomain", "api." + root, root, "", "", "", marker}, // under task root
|
||||
{"service", "www." + root, root, "https://www." + root + "/", "", "", marker}, // service under task root
|
||||
{"endpoint", "www." + root, root, "https://www." + root + "/a?" + marker + "=1", "GET", "", ""},
|
||||
{"subdomain", "dev." + srcRoot, srcRoot, "", "", "", marker}, // under source-task root
|
||||
{"endpoint", "198.51.100.9", "198.51.100.9", "http://198.51.100.9:8080/" + marker, "GET", "", ""}, // IP-literal host, ip col empty
|
||||
{"subdomain", "x." + out, out, "", "", "", marker}, // out of scope
|
||||
}
|
||||
var inScopeIDs, outIDs []int64
|
||||
for _, r := range rows {
|
||||
var id int64
|
||||
if err := d.QueryRow(`INSERT INTO assets(type,domain,root_domain,url,method,ip,page_title,task_ids,last_seen)
|
||||
VALUES ($1,NULLIF($2,''),NULLIF($3,''),NULLIF($4,''),NULLIF($5,''),NULLIF($6,''),NULLIF($7,''),ARRAY[$8]::bigint[],now())
|
||||
RETURNING id`,
|
||||
r.typ, r.domain, r.rootDom, r.url, r.method, r.ip, r.title, foreignTask).Scan(&id); err != nil {
|
||||
t.Fatalf("insert %s: %v", r.domain, err)
|
||||
}
|
||||
if r.rootDom == out {
|
||||
outIDs = append(outIDs, id)
|
||||
} else {
|
||||
inScopeIDs = append(inScopeIDs, id)
|
||||
}
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = d.Exec(`DELETE FROM assets WHERE id=ANY($1::bigint[]) OR id=ANY($2::bigint[])`, inScopeIDs, outIDs)
|
||||
})
|
||||
|
||||
got, err := assets.QueryDSLInScope(marker, "", task.ID, 50, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
gotIDs := map[int64]bool{}
|
||||
for _, a := range got {
|
||||
gotIDs[a.ID] = true
|
||||
}
|
||||
for _, id := range inScopeIDs {
|
||||
if !gotIDs[id] {
|
||||
t.Fatalf("in-scope asset %d missing from QueryDSLInScope result %v", id, gotIDs)
|
||||
}
|
||||
}
|
||||
for _, id := range outIDs {
|
||||
if gotIDs[id] {
|
||||
t.Fatalf("out-of-scope asset %d leaked into scoped query", id)
|
||||
}
|
||||
}
|
||||
|
||||
// GetByIDsInScope: in-scope id returns, out-of-scope id is dropped.
|
||||
mixed := append(append([]int64{}, inScopeIDs[0]), outIDs[0])
|
||||
byID, err := assets.GetByIDsInScope(task.ID, mixed)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(byID) != 1 || byID[0].ID != inScopeIDs[0] {
|
||||
t.Fatalf("GetByIDsInScope mixed ids → %+v, want only %d", byID, inScopeIDs[0])
|
||||
}
|
||||
|
||||
// taskID<=0 (non-task context) falls back to global: the out-of-scope row is reachable.
|
||||
global, err := assets.QueryDSLInScope(marker, "", 0, 50, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var sawOut bool
|
||||
for _, a := range global {
|
||||
if a.ID == outIDs[0] {
|
||||
sawOut = true
|
||||
}
|
||||
}
|
||||
if !sawOut {
|
||||
t.Fatalf("taskID<=0 should fall back to global and include out-of-scope asset %d", outIDs[0])
|
||||
}
|
||||
}
|
||||
@@ -52,6 +52,42 @@ tested AS (
|
||||
JOIN context_tasks ctx ON ctx.exploration_id=en.exploration_id
|
||||
)`
|
||||
|
||||
// scopeTargetCTE selects every asset that BELONGS to the current task's (and its
|
||||
// direct source tasks') declared scope — membership, not literal value: a
|
||||
// root_domain scope pulls in every subdomain / service / endpoint whose own
|
||||
// root_domain column equals it; an ip/cidr scope pulls in assets whose ip OR
|
||||
// IP-literal host falls inside the net. $1 is the task id. Unlike contextCoverageCTE
|
||||
// it carries neither the fact-anchor union nor the tested set — it is pure "in
|
||||
// declared scope", independent of what has already been touched. Used by the
|
||||
// agent's list_assets so a query returns the task's relevant assets, not the
|
||||
// whole shared库.
|
||||
const scopeTargetCTE = `
|
||||
context_tasks AS (
|
||||
SELECT t.id AS task_id
|
||||
FROM tasks t
|
||||
WHERE t.id=$1 AND t.deleted_at IS NULL
|
||||
UNION ALL
|
||||
SELECT source.id
|
||||
FROM task_relations relation
|
||||
JOIN tasks source ON source.id=relation.source_task_id AND source.deleted_at IS NULL
|
||||
WHERE relation.task_id=$1
|
||||
),
|
||||
target AS (
|
||||
SELECT DISTINCT a.id
|
||||
FROM assets a
|
||||
JOIN task_scope ts ON (
|
||||
(ts.kind='company' AND a.company_id = ts.company_id)
|
||||
OR (ts.kind='root_domain' AND a.root_domain = ts.domain)
|
||||
OR (ts.kind='subdomain' AND a.domain = ts.domain)
|
||||
OR (ts.kind IN ('ip','cidr') AND (ts.net >>= try_inet(a.ip) OR ts.net >>= try_inet(a.domain)))
|
||||
OR (ts.kind='icp' AND (
|
||||
lower(regexp_replace(COALESCE(a.icp,''), '[[:space:]]+', '', 'g')) = ts.value
|
||||
OR lower(regexp_replace(COALESCE(a.app_icp,''), '[[:space:]]+', '', 'g')) = ts.value
|
||||
))
|
||||
)
|
||||
JOIN context_tasks ctx ON ctx.task_id=ts.task_id
|
||||
)`
|
||||
|
||||
// ListTaskScopeWithSources returns the current task's scope followed by the
|
||||
// scopes of its direct source tasks. TaskScope.TaskID preserves provenance.
|
||||
func (s *AssetStore) ListTaskScopeWithSources(taskID int64) ([]TaskScope, error) {
|
||||
|
||||
Reference in new issue
Block a user