forge-tools-ssh/internal/tools/core.go

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
}