opsxdb/server/internal/aitools/runtime_test.go

package aitools

import (
	"context"
	"encoding/base64"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"opsxdb/server/internal/db"
)

func TestIsReadOnlySQL(t *testing.T) {
	okCases := []string{
		"SELECT 1",
		"SHOW TABLES",
		"EXPLAIN SELECT * FROM t",
		"WITH c AS (SELECT 1 AS x) SELECT * FROM c",
		"PRAGMA table_info(t)",
		"DESCRIBE t",
	}
	for _, sql := range okCases {
		if ok, reason := IsReadOnlySQL(sql); !ok {
			t.Fatalf("expected read-only ok for %q: %s", sql, reason)
		}
	}
	bad := []string{
		"DELETE FROM t",
		"DELETE FROM t WHERE id=1",
		"DROP TABLE t",
		"INSERT INTO t VALUES (1)",
		"UPDATE t SET a=1",
		"CREATE TABLE t (id INT)",
		"ALTER TABLE t ADD c INT",
		"TRUNCATE TABLE t",
		"WITH c AS (SELECT 1) INSERT INTO t SELECT * FROM c",
	}
	for _, sql := range bad {
		if ok, _ := IsReadOnlySQL(sql); ok {
			t.Fatalf("expected reject for %q", sql)
		}
	}
}

func TestCheckRedisReadOnly(t *testing.T) {
	if _, err := CheckRedisReadOnly("GET foo"); err != nil {
		t.Fatal(err)
	}
	if _, err := CheckRedisReadOnly("TYPE foo"); err != nil {
		t.Fatal(err)
	}
	if _, err := CheckRedisReadOnly("SCAN 0 MATCH * COUNT 10"); err != nil {
		t.Fatal(err)
	}
	for _, cmd := range []string{"FLUSHDB", "FLUSHALL", "KEYS *", "SET a 1", "DEL a"} {
		if _, err := CheckRedisReadOnly(cmd); err == nil {
			t.Fatalf("expected deny for %q", cmd)
		}
	}
}

func TestWhitelistEmptyRejects(t *testing.T) {
	rt := NewRuntime(DefaultPolicy()) // empty whitelist
	rt.LoadConns = func() ([]SavedConnection, error) {
		return []SavedConnection{{ID: "c1", Name: "n", Driver: "sqlite"}}, nil
	}
	rt.FindConn = func(id string) (*SavedConnection, error) {
		return &SavedConnection{ID: id, Driver: "sqlite"}, nil
	}
	var audits []AuditEntry
	rt.Audit = func(e AuditEntry) { audits = append(audits, e) }

	_, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{
		"connection_id": "c1",
		"sql":           "SELECT 1",
	})
	if err == nil || !strings.Contains(err.Error(), "whitelist empty") {
		t.Fatalf("expected whitelist empty error, got %v", err)
	}
}

func TestDangerousSQLBlocked(t *testing.T) {
	rt := NewRuntime(Policy{
		AllowedConnectionIDs: []string{"c1"},
		Mode:                 ModeReadOnly,
		DefaultLimit:         100,
		MaxLimit:             500,
	})
	rt.FindConn = func(id string) (*SavedConnection, error) {
		return &SavedConnection{ID: id, Driver: "mysql", Host: "127.0.0.1", Port: 3306}, nil
	}
	called := false
	rt.ExecSQL = func(ctx context.Context, c db.ConnReq) db.ExecResult {
		called = true
		return db.ExecResult{}
	}
	_, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{
		"connection_id": "c1",
		"sql":           "DROP TABLE users",
	})
	if err == nil {
		t.Fatal("expected DROP blocked")
	}
	if called {
		t.Fatal("ExecSQL must not be called for dangerous/write SQL")
	}
}

func TestDecodeSecretOx1(t *testing.T) {
	plain := "p@ss词"
	enc := "ox1." + base64.StdEncoding.EncodeToString([]byte(plain))
	if got := DecodeSecret(enc); got != plain {
		t.Fatalf("got %q want %q", got, plain)
	}
	if DecodeSecret("plain") != "plain" {
		t.Fatal("passthrough")
	}
}

func TestSQLiteListAndExec(t *testing.T) {
	dir := t.TempDir()
	dbPath := filepath.Join(dir, "t.db")
	if err := os.WriteFile(dbPath, []byte{}, 0600); err != nil {
		t.Fatal(err)
	}
	init := db.ExecSQL(context.Background(), db.ConnReq{
		Driver: "sqlite", Database: dbPath,
		SQL: "CREATE TABLE person (id INTEGER PRIMARY KEY, name TEXT)",
	})
	if init.Error != "" {
		t.Fatal(init.Error)
	}
	_ = db.ExecSQL(context.Background(), db.ConnReq{
		Driver: "sqlite", Database: dbPath,
		SQL: "INSERT INTO person (name) VALUES ('a')", ConfirmDanger: true,
	})

	rt := NewRuntime(Policy{
		AllowedConnectionIDs: []string{"sqlite1"},
		Mode:                 ModeReadOnly,
		DefaultLimit:         100,
		MaxLimit:             500,
	})
	rt.FindConn = func(id string) (*SavedConnection, error) {
		return &SavedConnection{ID: id, Name: "local", Driver: "sqlite", Database: dbPath}, nil
	}
	rt.LoadConns = func() ([]SavedConnection, error) {
		return []SavedConnection{{ID: "sqlite1", Name: "local", Driver: "sqlite", Database: dbPath}}, nil
	}

	out, err := rt.Call(context.Background(), ToolListTables, map[string]any{"connection_id": "sqlite1"})
	if err != nil {
		t.Fatal(err)
	}
	m := out.(map[string]any)
	tables := m["tables"].([]string)
	found := false
	for _, tb := range tables {
		if tb == "person" {
			found = true
		}
	}
	if !found {
		t.Fatalf("tables=%v", tables)
	}

	desc, err := rt.Call(context.Background(), ToolDescribeTable, map[string]any{
		"connection_id": "sqlite1", "table": "person",
	})
	if err != nil {
		t.Fatal(err)
	}
	_ = desc

	q, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{
		"connection_id": "sqlite1",
		"sql":           "SELECT name FROM person",
		"limit":         10,
	})
	if err != nil {
		t.Fatal(err)
	}
	qm := q.(map[string]any)
	if qm["rowCount"].(int) < 1 {
		t.Fatalf("rowCount=%v", qm["rowCount"])
	}

	// write rejected
	if _, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{
		"connection_id": "sqlite1",
		"sql":           "INSERT INTO person (name) VALUES ('b')",
	}); err == nil {
		t.Fatal("insert must be rejected")
	}
}

func TestListConnectionsFiltersWhitelist(t *testing.T) {
	rt := NewRuntime(Policy{AllowedConnectionIDs: []string{"a"}, Mode: ModeReadOnly, DefaultLimit: 100, MaxLimit: 500})
	rt.LoadConns = func() ([]SavedConnection, error) {
		return []SavedConnection{
			{ID: "a", Name: "A", Driver: "mysql", Host: "h", Port: 3306, Password: "secret"},
			{ID: "b", Name: "B", Driver: "mysql", Host: "h2", Port: 3306},
		}, nil
	}
	out, err := rt.Call(context.Background(), ToolListConnections, nil)
	if err != nil {
		t.Fatal(err)
	}
	conns := out.(map[string]any)["connections"].([]ConnectionSummary)
	if len(conns) != 1 || conns[0].ID != "a" {
		t.Fatalf("%+v", conns)
	}
	// ensure no password field leaked via JSON
	raw := ResultJSON(out)
	if strings.Contains(raw, "secret") || strings.Contains(raw, "password") {
		t.Fatalf("password leaked: %s", raw)
	}
}

func TestToolNamesStable(t *testing.T) {
	names := ToolNames()
	if len(names) != 7 {
		t.Fatalf("want 7 tools, got %d", len(names))
	}
	if names[0] != ToolListConnections || names[len(names)-1] != ToolExecuteRedisCommand {
		t.Fatalf("%v", names)
	}
}

func TestExecuteQueryNeedConfirmInReadWrite(t *testing.T) {
	called := false
	rt := NewRuntime(Policy{
		Mode:                 ModeReadWrite,
		AllowedConnectionIDs: []string{"c1"},
		DefaultLimit:         100,
		MaxLimit:             500,
	})
	rt.FindConn = func(id string) (*SavedConnection, error) {
		return &SavedConnection{ID: id, Name: "d", Driver: "mysql", Host: "127.0.0.1", Port: 3306}, nil
	}
	rt.ExecSQL = func(_ context.Context, c db.ConnReq) db.ExecResult {
		called = true
		return db.ExecResult{RowCount: 1}
	}
	rt.Audit = func(AuditEntry) {}

	out, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{
		"connection_id": "c1",
		"sql":           "DELETE FROM users WHERE id=1",
	})
	if err != nil {
		t.Fatal(err)
	}
	m := out.(map[string]any)
	if m["needConfirm"] != true {
		t.Fatalf("want needConfirm, got %#v", m)
	}
	if called {
		t.Fatal("must not exec without confirm")
	}

	out2, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{
		"connection_id":  "c1",
		"sql":            "DELETE FROM users WHERE id=1",
		"confirm_danger": true,
	})
	if err != nil {
		t.Fatal(err)
	}
	if !called {
		t.Fatal("expected exec after confirm")
	}
	_ = out2
}

func TestExecuteQueryRejectConfirmWhenAgentFlag(t *testing.T) {
	called := false
	rt := NewRuntime(Policy{Mode: ModeReadWrite, AllowedConnectionIDs: []string{"c1"}, DefaultLimit: 100, MaxLimit: 500})
	rt.RejectCallerConfirm = true
	rt.FindConn = func(id string) (*SavedConnection, error) {
		return &SavedConnection{ID: id, Driver: "mysql", Host: "h", Port: 1}, nil
	}
	rt.ExecSQL = func(_ context.Context, c db.ConnReq) db.ExecResult {
		called = true
		return db.ExecResult{}
	}
	rt.Audit = func(AuditEntry) {}
	out, err := rt.Call(context.Background(), ToolExecuteQuery, map[string]any{
		"connection_id": "c1", "sql": "UPDATE t SET a=1", "confirm_danger": true,
	})
	if err != nil {
		t.Fatal(err)
	}
	if out.(map[string]any)["needConfirm"] != true {
		t.Fatalf("%#v", out)
	}
	if called {
		t.Fatal("agent must not auto-confirm")
	}
}