Initial commit: forge-tools-ssh — MCP-сервер для администрирования по SSH
This commit is contained in:
Maksim Totmin
@@ -0,0 +1,251 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
func registerDBTools(s *mcp.Server) {
|
||||
// db_query
|
||||
s.AddTool(&mcp.Tool{
|
||||
Name: "db_query",
|
||||
Description: "Execute SQL/CQL/MongoDB query inside a database container",
|
||||
InputSchema: schema(map[string]any{
|
||||
"container": strProps("Docker container name running the database", true),
|
||||
"db_type": strProps("Database engine type", true, "postgres", "mysql", "scylladb", "cassandra", "mongodb"),
|
||||
"query": strProps("Query to execute", true),
|
||||
"database": strProps("Database/keyspace name", false),
|
||||
"username": strProps("Database username", false),
|
||||
"password": strProps("Database password", false),
|
||||
"timeout": intProps("Query timeout in seconds (default: 60)", false),
|
||||
"target": strProps("Connection alias (default: primary)", false),
|
||||
}, []string{"container", "db_type", "query"}),
|
||||
}, dbQueryHandler)
|
||||
|
||||
// db_schema
|
||||
s.AddTool(&mcp.Tool{
|
||||
Name: "db_schema",
|
||||
Description: "Get database schema (tables/collections list)",
|
||||
InputSchema: schema(map[string]any{
|
||||
"container": strProps("Docker container name", true),
|
||||
"db_type": strProps("Database engine type", true, "postgres", "mysql", "scylladb", "cassandra", "mongodb"),
|
||||
"database": strProps("Database/keyspace name", false),
|
||||
"username": strProps("Database username", false),
|
||||
"password": strProps("Database password", false),
|
||||
"target": strProps("Connection alias (default: primary)", false),
|
||||
}, []string{"container", "db_type"}),
|
||||
}, dbSchemaHandler)
|
||||
|
||||
// list_db_containers
|
||||
s.AddTool(&mcp.Tool{
|
||||
Name: "list_db_containers",
|
||||
Description: "Find Docker containers that look like databases",
|
||||
InputSchema: schema(map[string]any{
|
||||
"target": strProps("Connection alias (default: primary)", false),
|
||||
}, nil),
|
||||
}, listDBContainersHandler)
|
||||
}
|
||||
|
||||
func dbQueryHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
m := manager()
|
||||
if m == nil {
|
||||
return errorResult("no SSH manager initialized"), nil
|
||||
}
|
||||
args := requestArgs(req)
|
||||
container, _ := requireString(args, "container")
|
||||
dbType, _ := requireString(args, "db_type")
|
||||
query, _ := requireString(args, "query")
|
||||
database := getString(args, "database", "")
|
||||
username := getString(args, "username", "")
|
||||
password := getString(args, "password", "")
|
||||
timeout := getInt(args, "timeout", 60)
|
||||
target := getString(args, "target", "primary")
|
||||
|
||||
if err := checkDockerAvailable(ctx, m, target); err != nil {
|
||||
return errorResult(err.Error()), nil
|
||||
}
|
||||
|
||||
var cmd string
|
||||
switch dbType {
|
||||
case "postgres":
|
||||
user := "postgres"
|
||||
if username != "" {
|
||||
user = username
|
||||
}
|
||||
db := database
|
||||
if db == "" {
|
||||
db = "postgres"
|
||||
}
|
||||
cmd = fmt.Sprintf("docker exec -e PGPASSWORD=%s %s timeout %d psql -U %s -d %s -c %s 2>&1",
|
||||
shellQuote(password), shellQuote(container), timeout, shellQuote(user), shellQuote(db), shellQuote(query))
|
||||
|
||||
case "mysql":
|
||||
user := "root"
|
||||
if username != "" {
|
||||
user = username
|
||||
}
|
||||
cmd = fmt.Sprintf("docker exec %s timeout %d mysql -u%s", shellQuote(container), timeout, shellQuote(user))
|
||||
if password != "" {
|
||||
cmd += fmt.Sprintf(" -p%s", shellQuote(password))
|
||||
}
|
||||
if database != "" {
|
||||
cmd += fmt.Sprintf(" %s", shellQuote(database))
|
||||
}
|
||||
cmd += fmt.Sprintf(" -e %s 2>&1", shellQuote(query))
|
||||
|
||||
case "scylladb", "cassandra":
|
||||
cmd = fmt.Sprintf("docker exec %s timeout %d cqlsh", shellQuote(container), timeout)
|
||||
if username != "" {
|
||||
cmd += fmt.Sprintf(" -u %s", shellQuote(username))
|
||||
}
|
||||
if password != "" {
|
||||
cmd += fmt.Sprintf(" -p %s", shellQuote(password))
|
||||
}
|
||||
cmd += fmt.Sprintf(" -e %s 2>&1", shellQuote(query))
|
||||
|
||||
case "mongodb":
|
||||
db := database
|
||||
if db == "" {
|
||||
db = "admin"
|
||||
}
|
||||
cmd = fmt.Sprintf("docker exec %s timeout %d mongosh --quiet %s", shellQuote(container), timeout, shellQuote(db))
|
||||
if username != "" && password != "" {
|
||||
cmd += fmt.Sprintf(" -u %s -p %s --authenticationDatabase admin", shellQuote(username), shellQuote(password))
|
||||
}
|
||||
cmd += fmt.Sprintf(" --eval %s 2>&1", shellQuote(query))
|
||||
|
||||
default:
|
||||
return errorResult(fmt.Sprintf("Unsupported database type: %s. Supported: postgres, mysql, scylladb, cassandra, mongodb", dbType)), nil
|
||||
}
|
||||
|
||||
output, err := m.Execute(ctx, cmd, target)
|
||||
if err != nil {
|
||||
return errorResult(err.Error()), nil
|
||||
}
|
||||
return textResult(output), nil
|
||||
}
|
||||
|
||||
func dbSchemaHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
m := manager()
|
||||
if m == nil {
|
||||
return errorResult("no SSH manager initialized"), nil
|
||||
}
|
||||
args := requestArgs(req)
|
||||
container, _ := requireString(args, "container")
|
||||
dbType, _ := requireString(args, "db_type")
|
||||
database := getString(args, "database", "")
|
||||
username := getString(args, "username", "")
|
||||
password := getString(args, "password", "")
|
||||
target := getString(args, "target", "primary")
|
||||
|
||||
if err := checkDockerAvailable(ctx, m, target); err != nil {
|
||||
return errorResult(err.Error()), nil
|
||||
}
|
||||
|
||||
var query string
|
||||
switch dbType {
|
||||
case "postgres":
|
||||
query = "\\dt"
|
||||
case "mysql":
|
||||
query = "SHOW TABLES;"
|
||||
case "scylladb", "cassandra":
|
||||
if database != "" {
|
||||
query = fmt.Sprintf("DESCRIBE KEYSPACE %s;", database)
|
||||
} else {
|
||||
query = "DESCRIBE KEYSPACES;"
|
||||
}
|
||||
case "mongodb":
|
||||
query = "db.getCollectionNames()"
|
||||
default:
|
||||
return errorResult(fmt.Sprintf("Unsupported database type: %s", dbType)), nil
|
||||
}
|
||||
|
||||
var cmd string
|
||||
switch dbType {
|
||||
case "postgres":
|
||||
user := "postgres"
|
||||
if username != "" {
|
||||
user = username
|
||||
}
|
||||
db := database
|
||||
if db == "" {
|
||||
db = "postgres"
|
||||
}
|
||||
cmd = fmt.Sprintf("docker exec -e PGPASSWORD=%s %s psql -U %s -d %s -c %s 2>&1",
|
||||
shellQuote(password), shellQuote(container), shellQuote(user), shellQuote(db), shellQuote(query))
|
||||
case "mysql":
|
||||
user := "root"
|
||||
if username != "" {
|
||||
user = username
|
||||
}
|
||||
cmd = fmt.Sprintf("docker exec %s mysql -u%s", shellQuote(container), shellQuote(user))
|
||||
if password != "" {
|
||||
cmd += fmt.Sprintf(" -p%s", shellQuote(password))
|
||||
}
|
||||
if database != "" {
|
||||
cmd += fmt.Sprintf(" %s", shellQuote(database))
|
||||
}
|
||||
cmd += fmt.Sprintf(" -e %s 2>&1", shellQuote(query))
|
||||
case "scylladb", "cassandra":
|
||||
cmd = fmt.Sprintf("docker exec %s cqlsh", shellQuote(container))
|
||||
if username != "" {
|
||||
cmd += fmt.Sprintf(" -u %s", shellQuote(username))
|
||||
}
|
||||
if password != "" {
|
||||
cmd += fmt.Sprintf(" -p %s", shellQuote(password))
|
||||
}
|
||||
cmd += fmt.Sprintf(" -e %s 2>&1", shellQuote(query))
|
||||
case "mongodb":
|
||||
db := database
|
||||
if db == "" {
|
||||
db = "admin"
|
||||
}
|
||||
cmd = fmt.Sprintf("docker exec %s mongosh --quiet %s", shellQuote(container), shellQuote(db))
|
||||
if username != "" && password != "" {
|
||||
cmd += fmt.Sprintf(" -u %s -p %s --authenticationDatabase admin", shellQuote(username), shellQuote(password))
|
||||
}
|
||||
cmd += fmt.Sprintf(" --eval %s 2>&1", shellQuote(query))
|
||||
}
|
||||
|
||||
output, err := m.Execute(ctx, cmd, target)
|
||||
if err != nil {
|
||||
return errorResult(err.Error()), nil
|
||||
}
|
||||
return textResult(output), nil
|
||||
}
|
||||
|
||||
func listDBContainersHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
m := manager()
|
||||
if m == nil {
|
||||
return errorResult("no SSH manager initialized"), nil
|
||||
}
|
||||
target := getString(requestArgs(req), "target", "primary")
|
||||
|
||||
if err := checkDockerAvailable(ctx, m, target); err != nil {
|
||||
return errorResult(err.Error()), nil
|
||||
}
|
||||
|
||||
cmd := `docker ps --format '{{.Names}}|{{.Image}}' | while read line; do
|
||||
name=$(echo "$line" | cut -d'|' -f1)
|
||||
image=$(echo "$line" | cut -d'|' -f2)
|
||||
case "$image" in
|
||||
*postgres*) echo "$name|$image|postgres" ;;
|
||||
*mysql*|*mariadb*) echo "$name|$image|mysql" ;;
|
||||
*scylla*) echo "$name|$image|scylladb" ;;
|
||||
*cassandra*) echo "$name|$image|cassandra" ;;
|
||||
*mongo*) echo "$name|$image|mongodb" ;;
|
||||
*redis*) echo "$name|$image|redis" ;;
|
||||
esac
|
||||
done 2>/dev/null`
|
||||
|
||||
output, err := m.Execute(ctx, cmd, target)
|
||||
if err != nil {
|
||||
return errorResult(err.Error()), nil
|
||||
}
|
||||
if trimOutput(output) == "" {
|
||||
return textResult("No database containers found"), nil
|
||||
}
|
||||
return textResult(output), nil
|
||||
}
|
||||
Reference in New Issue
Block a user