Initial commit: forge-tools-ssh — MCP-сервер для администрирования по SSH
This commit is contained in:
Maksim Totmin
@@ -0,0 +1,186 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"forge-tools-ssh/internal/ssh"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
var mgr *ssh.Manager
|
||||
|
||||
// SetManager binds the single SSH Manager for the process.
|
||||
// forge держит 1 MCP-подпроцесс/агента (stdio), поэтому менеджер один.
|
||||
func SetManager(m *ssh.Manager) {
|
||||
mgr = m
|
||||
}
|
||||
|
||||
// manager возвращает одиночный SSH Manager процесса.
|
||||
func manager() *ssh.Manager {
|
||||
return mgr
|
||||
}
|
||||
|
||||
// registerCoreTools регистрирует базовые SSH-инструменты.
|
||||
func registerCoreTools(s *mcp.Server) {
|
||||
// connect
|
||||
s.AddTool(&mcp.Tool{
|
||||
Name: "connect",
|
||||
Description: "Establish an SSH connection to a remote host. Используй profile (имя из ssh.json), чтобы подключиться по предопределённому профилю; либо host+username для ad-hoc (host должен пройти allowlist политики).",
|
||||
InputSchema: schema(map[string]any{
|
||||
"profile": strProps("Имя профиля из per-agent ssh.json (креды задаёт оператор)", false),
|
||||
"host": strProps("Hostname or IP address", false),
|
||||
"username": strProps("SSH username", false),
|
||||
"port": intProps("SSH port (default: 22)", false),
|
||||
"password": strProps("SSH password (optional if using key)", false),
|
||||
"private_key_path": strProps("Path to private key file", false),
|
||||
"alias": strProps("Connection alias (auto-generated if not provided)", false),
|
||||
"via": strProps("Jump host alias for tunneling", false),
|
||||
"target": strProps("Optional destination host routed through a PAM gateway (SafeInspect). The gateway receives 'username@target' and proxies the session.", false),
|
||||
"target_password": strProps("Password of the target server account, requested by the PAM gateway as a second authentication stage", false),
|
||||
}, nil),
|
||||
}, connectHandler)
|
||||
|
||||
// disconnect
|
||||
s.AddTool(&mcp.Tool{
|
||||
Name: "disconnect",
|
||||
Description: "Close an SSH connection",
|
||||
InputSchema: schema(map[string]any{
|
||||
"alias": strProps("Connection alias to disconnect (all if empty)", false),
|
||||
}, nil),
|
||||
}, disconnectHandler)
|
||||
|
||||
// run
|
||||
registerPatternTool(s, &mcp.Tool{
|
||||
Name: "run",
|
||||
Description: "Execute a shell command on the remote host. Use timeout for long-running tasks.",
|
||||
InputSchema: schema(map[string]any{
|
||||
"command": strProps("Shell command to execute", true),
|
||||
"target": strProps("Connection alias (default: primary)", false),
|
||||
"timeout": intProps("Command timeout in seconds (default: 120)", false),
|
||||
}, []string{"command"}),
|
||||
}, runPatternsFn, runHandler)
|
||||
|
||||
// identity
|
||||
s.AddTool(&mcp.Tool{
|
||||
Name: "identity",
|
||||
Description: "Get the server's public SSH key for authorized_keys",
|
||||
InputSchema: schema(nil, nil),
|
||||
}, identityHandler)
|
||||
|
||||
// info
|
||||
s.AddTool(&mcp.Tool{
|
||||
Name: "info",
|
||||
Description: "Get remote system information (OS, kernel, hostname)",
|
||||
InputSchema: schema(map[string]any{
|
||||
"target": strProps("Connection alias (default: primary)", false),
|
||||
}, nil),
|
||||
}, infoHandler)
|
||||
}
|
||||
|
||||
func connectHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
m := manager()
|
||||
if m == nil {
|
||||
return errorResult("no SSH manager initialized"), nil
|
||||
}
|
||||
|
||||
args := requestArgs(req)
|
||||
opts := ssh.ConnectOptions{
|
||||
Host: getString(args, "host", ""),
|
||||
Username: getString(args, "username", ""),
|
||||
Port: getInt(args, "port", 22),
|
||||
Password: getString(args, "password", ""),
|
||||
PrivateKeyPath: getString(args, "private_key_path", ""),
|
||||
Alias: getString(args, "alias", ""),
|
||||
Via: getString(args, "via", ""),
|
||||
Target: getString(args, "target", ""),
|
||||
TargetPassword: getString(args, "target_password", ""),
|
||||
Profile: getString(args, "profile", ""),
|
||||
}
|
||||
|
||||
resultAlias, err := m.Connect(ctx, opts)
|
||||
if err != nil {
|
||||
return errorResult(err.Error()), nil
|
||||
}
|
||||
|
||||
dest := opts.Host
|
||||
if opts.Target != "" {
|
||||
dest = fmt.Sprintf("%s via PAM -> %s", opts.Host, opts.Target)
|
||||
}
|
||||
return textResult(fmt.Sprintf("Connected to %s@%s (alias: %s)", opts.Username, dest, resultAlias)), nil
|
||||
}
|
||||
|
||||
func disconnectHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
m := manager()
|
||||
if m == nil {
|
||||
return errorResult("no SSH manager initialized"), nil
|
||||
}
|
||||
|
||||
alias := getString(requestArgs(req), "alias", "")
|
||||
msg, err := m.Disconnect(alias)
|
||||
if err != nil {
|
||||
return errorResult(err.Error()), nil
|
||||
}
|
||||
return textResult(msg), nil
|
||||
}
|
||||
|
||||
func runHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
m := manager()
|
||||
if m == nil {
|
||||
return errorResult("no SSH manager initialized"), nil
|
||||
}
|
||||
|
||||
args := requestArgs(req)
|
||||
command := getString(args, "command", "")
|
||||
if command == "" {
|
||||
return errorResult("command is required"), nil
|
||||
}
|
||||
target := getString(args, "target", "primary")
|
||||
timeout := getInt(args, "timeout", 120)
|
||||
|
||||
if timeout > 0 {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
output, err := m.Execute(ctx, command, target)
|
||||
if err != nil {
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
return errorResult(fmt.Sprintf("Command timed out after %ds", timeout)), nil
|
||||
}
|
||||
return errorResult(err.Error()), nil
|
||||
}
|
||||
return textResult(output), nil
|
||||
}
|
||||
|
||||
func identityHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
m := manager()
|
||||
if m == nil {
|
||||
return errorResult("no SSH manager initialized"), nil
|
||||
}
|
||||
|
||||
pubKey, err := m.GetPublicKey()
|
||||
if err != nil {
|
||||
return errorResult(err.Error()), nil
|
||||
}
|
||||
formatted := fmt.Sprintf("SSH Public Key:\n\n```\n%s```\n\nAdd this to ~/.ssh/authorized_keys on remote servers.", pubKey)
|
||||
return textResult(formatted), nil
|
||||
}
|
||||
|
||||
func infoHandler(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")
|
||||
cmd := `echo "Hostname: $(hostname)"; echo "OS: $(cat /etc/os-release 2>/dev/null | grep PRETTY_NAME | cut -d'"' -f2 || uname -s)"; echo "Kernel: $(uname -r)"; echo "Arch: $(uname -m)"; echo "Shell: $SHELL"`
|
||||
output, err := m.Execute(ctx, cmd, target)
|
||||
if err != nil {
|
||||
return errorResult(err.Error()), nil
|
||||
}
|
||||
return textResult(output), nil
|
||||
}
|
||||
Reference in New Issue
Block a user