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 WHERE id=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") } }