package db import ( "fmt" "strconv" "strings" "unicode" ) // Expr is one leaf DSL clause. type Expr struct { Field string // empty = bare-text full-text search Op string // "=", "==", "!=", ">", ">=", "<", "<=" Value string } // astNode is a node in the parsed DSL expression tree. type astNode struct { kind string // "and", "or", "leaf" children []*astNode expr *Expr // only for "leaf" } func andNode(cs []*astNode) *astNode { return &astNode{kind: "and", children: cs} } func orNode(cs []*astNode) *astNode { return &astNode{kind: "or", children: cs} } func leafNode(e Expr) *astNode { return &astNode{kind: "leaf", expr: &e} } // knownStringFields maps DSL field name → SQL column name. // NOTE: "type" is intentionally excluded — it is a separate parameter, not a DSL field. var knownStringFields = map[string]string{ "domain": "domain", "root_domain": "root_domain", "ip": "ip", "url": "url", "page_title": "page_title", "title": "page_title", "icp": "icp", "service_name": "service_name", "app_name": "app_name", "bundle_id": "bundle_id", "category": "category", "app_icp": "app_icp", "method": "method", "service_type": "service_type", "record_type": "record_type", } // knownArrayFields maps DSL field name → SQL column name (array). var knownArrayFields = map[string]string{ "technology": "technologies", "technologies": "technologies", "tech": "technologies", } // knownNumericFields maps DSL field name → SQL column name (integer). var knownNumericFields = map[string]string{ "port": "port", "status_code": "status_code", "status": "status_code", } func isKnownField(f string) bool { f = strings.ToLower(f) _, s := knownStringFields[f] _, a := knownArrayFields[f] _, n := knownNumericFields[f] return s || a || n || f == "company_id" || f == "task_id" } // ── tokeniser ──────────────────────────────────────────────────────────────── const ( tkField = "FIELD" tkBare = "BARE" tkAnd = "AND" tkOr = "OR" tkLP = "LPAREN" tkRP = "RPAREN" tkEOF = "EOF" ) type tok struct { kind string expr *Expr // set for tkField and tkBare } func tokenize(s string) ([]tok, error) { var tokens []tok i := 0 for i < len(s) { for i < len(s) && unicode.IsSpace(rune(s[i])) { i++ } if i >= len(s) { break } switch s[i] { case '(': tokens = append(tokens, tok{kind: tkLP}) i++ case ')': tokens = append(tokens, tok{kind: tkRP}) i++ default: if expr, end, ok := tryParseFieldExpr(s, i); ok { tokens = append(tokens, tok{kind: tkField, expr: &expr}) i = end continue } word, end := readToken(s, i) if word == "" { i++ continue } switch strings.ToUpper(word) { case "AND": tokens = append(tokens, tok{kind: tkAnd}) case "OR": tokens = append(tokens, tok{kind: tkOr}) default: e := Expr{Field: "", Op: "=", Value: word} tokens = append(tokens, tok{kind: tkBare, expr: &e}) } i = end } } tokens = append(tokens, tok{kind: tkEOF}) return tokens, nil } // ── parser ─────────────────────────────────────────────────────────────────── // // Grammar (AND binds tighter than OR): // expr = or_expr // or_expr = and_expr (OR and_expr)* // and_expr = atom (AND atom)* // atom = FIELD | BARE | '(' expr ')' type dslParser struct { tokens []tok pos int } func (p *dslParser) peek() tok { if p.pos >= len(p.tokens) { return tok{kind: tkEOF} } return p.tokens[p.pos] } func (p *dslParser) consume() tok { t := p.peek() p.pos++ return t } func (p *dslParser) parseOr() (*astNode, error) { left, err := p.parseAnd() if err != nil { return nil, err } children := []*astNode{left} for p.peek().kind == tkOr { p.consume() right, err := p.parseAnd() if err != nil { return nil, err } children = append(children, right) } if len(children) == 1 { return children[0], nil } return orNode(children), nil } func (p *dslParser) parseAnd() (*astNode, error) { left, err := p.parseAtom() if err != nil { return nil, err } children := []*astNode{left} for p.peek().kind == tkAnd { p.consume() right, err := p.parseAtom() if err != nil { return nil, err } children = append(children, right) } if len(children) == 1 { return children[0], nil } return andNode(children), nil } func (p *dslParser) parseAtom() (*astNode, error) { t := p.peek() switch t.kind { case tkField, tkBare: p.consume() return leafNode(*t.expr), nil case tkLP: p.consume() node, err := p.parseOr() if err != nil { return nil, err } if p.peek().kind != tkRP { return nil, fmt.Errorf("DSL 语法错误:缺少右括号 ')'") } p.consume() return node, nil case tkEOF: return nil, fmt.Errorf("DSL 语法错误:表达式不完整") default: return nil, fmt.Errorf("DSL 语法错误:意外的 token '%s'", t.kind) } } // ParseDSL parses a DSL query string into an expression tree. // // Syntax: // // field=value fuzzy match (ILIKE '%value%') // field==value exact match // field!=value exclude fuzzy // port>8080 numeric comparison // bare word full-text fuzzy across all main text fields // // Operators: AND OR (case-insensitive), parentheses for grouping. // AND binds tighter than OR. func ParseDSL(s string) (*astNode, error) { if strings.TrimSpace(s) == "" { return nil, nil } tokens, err := tokenize(s) if err != nil { return nil, err } p := &dslParser{tokens: tokens} node, err := p.parseOr() if err != nil { return nil, err } if p.peek().kind != tkEOF { return nil, fmt.Errorf("DSL 语法错误:意外的内容 '%s'", p.peek().kind) } return node, nil } // ── SQL builder ────────────────────────────────────────────────────────────── // fullTextCols are searched for bare-text tokens. var fullTextCols = []string{ "domain", "root_domain", "ip", "url", "page_title", "icp", "service_name", "app_name", "app_description", } 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", b.base+len(b.args)) } func (b *whereBuilder) build(node *astNode) (string, error) { switch node.kind { case "and": parts := make([]string, 0, len(node.children)) for _, child := range node.children { clause, err := b.build(child) if err != nil { return "", err } parts = append(parts, "("+clause+")") } return strings.Join(parts, " AND "), nil case "or": parts := make([]string, 0, len(node.children)) for _, child := range node.children { clause, err := b.build(child) if err != nil { return "", err } parts = append(parts, "("+clause+")") } return strings.Join(parts, " OR "), nil case "leaf": return b.buildLeaf(*node.expr) } return "", fmt.Errorf("unknown node kind: %s", node.kind) } func (b *whereBuilder) buildLeaf(e Expr) (string, error) { f := strings.ToLower(e.Field) // bare-text: OR across all text fields + arrays if f == "" { p := b.next("%" + e.Value + "%") var parts []string for _, col := range fullTextCols { parts = append(parts, col+" ILIKE "+p) } parts = append(parts, "EXISTS (SELECT 1 FROM unnest(technologies) t(v) WHERE v ILIKE "+p+")", "EXISTS (SELECT 1 FROM unnest(bound_domains) t(v) WHERE v ILIKE "+p+")", ) return "(" + strings.Join(parts, " OR ") + ")", nil } // task_id: $N = ANY(task_ids) if f == "task_id" { n, err := strconv.ParseInt(e.Value, 10, 64) if err != nil { return "", fmt.Errorf("task_id 需要整数值: %s", e.Value) } return b.next(n) + " = ANY(task_ids)", nil } // company_id: exact integer if f == "company_id" { n, err := strconv.ParseInt(e.Value, 10, 64) if err != nil { return "", fmt.Errorf("company_id 需要整数值: %s", e.Value) } return "company_id = " + b.next(n), nil } // numeric fields if col, ok := knownNumericFields[f]; ok { n, err := strconv.Atoi(e.Value) if err != nil { return "", fmt.Errorf("字段 %s 需要整数值: %s", f, e.Value) } op := e.Op if op == "==" { op = "=" } if op != "=" && op != "!=" && op != ">" && op != ">=" && op != "<" && op != "<=" { return "", fmt.Errorf("字段 %s 不支持运算符 %s", f, e.Op) } return fmt.Sprintf("%s %s %s", col, op, b.next(n)), nil } // array fields if col, ok := knownArrayFields[f]; ok { switch e.Op { case "==": return b.next(e.Value) + " = ANY(" + col + ")", nil case "!=": return "NOT (" + b.next(e.Value) + " = ANY(" + col + "))", nil case "=": p := b.next("%" + e.Value + "%") return "EXISTS (SELECT 1 FROM unnest(" + col + ") t(v) WHERE v ILIKE " + p + ")", nil default: return "", fmt.Errorf("数组字段 %s 不支持运算符 %s", f, e.Op) } } // string fields if col, ok := knownStringFields[f]; ok { switch e.Op { case "=": return col + " ILIKE " + b.next("%"+e.Value+"%"), nil case "==": return col + " = " + b.next(e.Value), nil case "!=": return col + " NOT ILIKE " + b.next("%"+e.Value+"%"), nil default: return "", fmt.Errorf("字符串字段 %s 不支持运算符 %s", f, e.Op) } } return "", fmt.Errorf("未知字段: %s", f) } 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{base: base} clause, err := b.build(node) if err != nil { return "", nil, err } return clause, b.args, nil } // ── helpers (shared with parser) ───────────────────────────────────────────── // tryParseFieldExpr tries to parse "field op value" at pos. func tryParseFieldExpr(s string, pos int) (Expr, int, bool) { i := pos if i >= len(s) || !isIdentStart(s[i]) { return Expr{}, pos, false } for i < len(s) && isIdentChar(s[i]) { i++ } field := strings.ToLower(s[pos:i]) if !isKnownField(field) { return Expr{}, pos, false } if i >= len(s) { return Expr{}, pos, false } var op string switch { case i+1 < len(s) && (s[i] == '=' || s[i] == '!' || s[i] == '>' || s[i] == '<') && s[i+1] == '=': op = s[i : i+2] i += 2 case s[i] == '>' || s[i] == '<' || s[i] == '=': op = string(s[i]) i++ default: return Expr{}, pos, false } value, end := readToken(s, i) if end == i { return Expr{}, pos, false } return Expr{Field: field, Op: op, Value: value}, end, true } // readToken reads a quoted or unquoted token starting at pos. func readToken(s string, pos int) (string, int) { if pos >= len(s) { return "", pos } if s[pos] == '"' { i := pos + 1 for i < len(s) && s[i] != '"' { i++ } val := s[pos+1 : i] if i < len(s) { i++ } return val, i } i := pos for i < len(s) && !unicode.IsSpace(rune(s[i])) && s[i] != '(' && s[i] != ')' { i++ } return s[pos:i], i } func isIdentStart(c byte) bool { return c == '_' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') } func isIdentChar(c byte) bool { return isIdentStart(c) || (c >= '0' && c <= '9') } // ── QueryDSL ───────────────────────────────────────────────────────────────── // ValidateDSL parses and compiles a DSL expression without touching the // database. HTTP callers use it to distinguish client syntax errors from query // failures, which must remain server errors. func ValidateDSL(dsl string) error { node, err := ParseDSL(dsl) if err != nil { return err } _, _, err = buildDSLWhere(node) return err } // CountDSL returns the total number of assets matching a DSL expression (and optional // type), for server-side pagination — same WHERE as QueryDSL, without LIMIT/OFFSET. // taskID > 0 scopes the count to assets attached to that task. func (s *AssetStore) CountDSL(dsl, typ string, taskID int64) (int, error) { node, err := ParseDSL(dsl) if err != nil { return 0, err } where, args, err := buildDSLWhere(node) if err != nil { return 0, err } if typ != "" { args = append(args, typ) where += fmt.Sprintf(" AND type = $%d", len(args)) } if taskID > 0 { args = append(args, taskID) where += fmt.Sprintf(" AND $%d = ANY(task_ids)", len(args)) } var n int err = s.db.QueryRow("SELECT count(*) FROM assets WHERE "+where, args...).Scan(&n) return n, err } // QueryDSL executes a DSL query string against the asset store. // typ is an optional asset type filter applied independently of the DSL expression. // taskID > 0 scopes results to assets attached to that task and hydrates each // row's per-task source metadata (as QueryByTask does). func (s *AssetStore) QueryDSL(dsl, typ string, taskID int64, limit, offset int) ([]*Asset, error) { if limit <= 0 { limit = 50 } if offset < 0 { offset = 0 } node, err := ParseDSL(dsl) if err != nil { return nil, err } where, args, err := buildDSLWhere(node) if err != nil { return nil, err } if typ != "" { args = append(args, typ) where += fmt.Sprintf(" AND type = $%d", len(args)) } if taskID > 0 { args = append(args, taskID) where += fmt.Sprintf(" AND $%d = ANY(task_ids)", len(args)) } args = append(args, limit, offset) q := 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 taskID > 0 { if err := s.hydrateTaskAssetSources(taskID, assets); err != nil { return nil, err } } 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 }