First Commit

db/asset_dsl.go

@@ -0,0 +1,637 @@
+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
+}