105 lines
2.9 KiB
Go
105 lines
2.9 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"strings"
|
|
|
|
"opsxdb/server/internal/aitools"
|
|
)
|
|
|
|
func must(ok bool, msg string) {
|
|
if !ok {
|
|
fmt.Println("FAIL:", msg)
|
|
os.Exit(1)
|
|
}
|
|
fmt.Println("OK:", msg)
|
|
}
|
|
|
|
func dump(v any) string {
|
|
b, _ := json.MarshalIndent(v, "", " ")
|
|
return string(b)
|
|
}
|
|
|
|
func main() {
|
|
rt, err := aitools.DefaultRuntime()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
mysqlID := "accept-mysql-demo"
|
|
redisID := "accept-redis-demo"
|
|
ctx := context.Background()
|
|
|
|
out, err := rt.Call(ctx, aitools.ToolListConnections, map[string]any{})
|
|
s := dump(out)
|
|
must(err == nil && strings.Contains(s, mysqlID) && !strings.Contains(s, `"password"`), "list_connections")
|
|
|
|
out, err = rt.Call(ctx, aitools.ToolListDatabases, map[string]any{"connection_id": mysqlID})
|
|
s = dump(out)
|
|
must(err == nil && strings.Contains(s, "demo"), "list_databases: "+fmt.Sprint(err)+" "+s)
|
|
|
|
out, err = rt.Call(ctx, aitools.ToolListTables, map[string]any{"connection_id": mysqlID, "database": "demo"})
|
|
s = dump(out)
|
|
must(err == nil, "list_tables: "+fmt.Sprint(err)+" "+s)
|
|
fmt.Println(s[:min(600, len(s))])
|
|
|
|
var tables []string
|
|
if m, ok := out.(map[string]any); ok {
|
|
if arr, ok := m["tables"].([]any); ok {
|
|
for _, x := range arr {
|
|
tables = append(tables, fmt.Sprint(x))
|
|
}
|
|
}
|
|
}
|
|
if len(tables) > 0 {
|
|
t := tables[0]
|
|
_, err = rt.Call(ctx, aitools.ToolDescribeTable, map[string]any{"connection_id": mysqlID, "database": "demo", "table": t})
|
|
must(err == nil, "describe_table "+t+": "+fmt.Sprint(err))
|
|
out, err = rt.Call(ctx, aitools.ToolGetSchemaContext, map[string]any{"connection_id": mysqlID, "database": "demo"})
|
|
must(err == nil && len(dump(out)) > 20, "get_schema_context")
|
|
_ = out
|
|
}
|
|
|
|
out, err = rt.Call(ctx, aitools.ToolExecuteQuery, map[string]any{
|
|
"connection_id": mysqlID,
|
|
"database": "demo",
|
|
"sql": "SELECT 1 AS n, DATABASE() AS db",
|
|
})
|
|
s = dump(out)
|
|
must(err == nil && strings.Contains(s, "demo"), "execute_query SELECT: "+fmt.Sprint(err)+" "+s)
|
|
fmt.Println(s[:min(500, len(s))])
|
|
|
|
_, err = rt.Call(ctx, aitools.ToolExecuteQuery, map[string]any{
|
|
"connection_id": mysqlID,
|
|
"sql": "DROP TABLE IF EXISTS __opsxdb_should_deny__",
|
|
})
|
|
must(err != nil, "DROP must fail, got nil")
|
|
|
|
out, err = rt.Call(ctx, aitools.ToolExecuteRedisCommand, map[string]any{
|
|
"connection_id": redisID,
|
|
"command": "PING",
|
|
})
|
|
s = dump(out)
|
|
must(err == nil && strings.Contains(strings.ToUpper(s), "PONG"), "redis PING: "+fmt.Sprint(err)+" "+s)
|
|
|
|
_, err = rt.Call(ctx, aitools.ToolExecuteRedisCommand, map[string]any{
|
|
"connection_id": redisID,
|
|
"command": "FLUSHDB",
|
|
})
|
|
must(err != nil, "FLUSHDB must fail")
|
|
|
|
_, err = rt.Call(ctx, aitools.ToolListTables, map[string]any{"connection_id": "not-allowed"})
|
|
must(err != nil, "whitelist reject")
|
|
|
|
fmt.Println("ALL ACCEPTANCE PASSED")
|
|
}
|
|
|
|
func min(a, b int) int {
|
|
if a < b {
|
|
return a
|
|
}
|
|
return b
|
|
}
|