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

package tools

import (
	"context"
	"fmt"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

func registerNetworkTools(s *mcp.Server) {
	// net_stat
	s.AddTool(&mcp.Tool{
		Name:        "net_stat",
		Description: "Check listening ports (ss/netstat)",
		InputSchema: schema(map[string]any{
			"port":   intProps("Filter by specific port", false),
			"target": strProps("Connection alias (default: primary)", false),
		}, nil),
	}, netStatHandler)

	// search_files
	s.AddTool(&mcp.Tool{
		Name:        "search_files",
		Description: "Find files using POSIX find",
		InputSchema: schema(map[string]any{
			"pattern":   strProps("File name pattern (supports wildcards)", true),
			"path":      strProps("Search path (default: /)", false),
			"max_depth": intProps("Maximum directory depth", false),
			"type":      strProps("Filter by type: f (file), d (directory)", false, "f", "d"),
			"target":    strProps("Connection alias (default: primary)", false),
		}, []string{"pattern"}),
	}, searchFilesHandler)

	// search_text
	s.AddTool(&mcp.Tool{
		Name:        "search_text",
		Description: "Search text in files using grep",
		InputSchema: schema(map[string]any{
			"pattern":     strProps("Search pattern", true),
			"path":        strProps("File or directory path", true),
			"recursive":   boolProps("Search recursively", false),
			"ignore_case": boolProps("Case-insensitive search", false),
			"context":     intProps("Lines of context around matches", false),
			"target":      strProps("Connection alias (default: primary)", false),
		}, []string{"pattern", "path"}),
	}, searchTextHandler)

	// package_manage
	s.AddTool(&mcp.Tool{
		Name:        "package_manage",
		Description: "Install/remove/check packages (apt, apk, dnf, yum)",
		InputSchema: schema(map[string]any{
			"action":  strProps("Package management action", true, "install", "remove", "check", "list"),
			"package": strProps("Package name (required for install/remove/check)", false),
			"target":  strProps("Connection alias (default: primary)", false),
		}, []string{"action"}),
	}, packageManageHandler)
}

func netStatHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	m := manager()
	if m == nil {
		return errorResult("no SSH manager initialized"), nil
	}
	args := requestArgs(req)
	port := getInt(args, "port", 0)
	target := getString(args, "target", "primary")

	var cmd string
	if port > 0 {
		cmd = fmt.Sprintf("ss -tlnp 2>/dev/null | grep ':%d ' || netstat -an 2>/dev/null | grep -i listen | grep '[\\.: ]%d '", port, port)
	} else {
		cmd = "ss -tlnp 2>/dev/null || netstat -an 2>/dev/null | grep -i listen"
	}
	output, err := m.Execute(ctx, cmd, target)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	return textResult(output), nil
}

func searchFilesHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	m := manager()
	if m == nil {
		return errorResult("no SSH manager initialized"), nil
	}
	args := requestArgs(req)
	pattern, _ := requireString(args, "pattern")
	path := getString(args, "path", "/")
	maxDepth := getInt(args, "max_depth", 0)
	fileType := getString(args, "type", "")
	target := getString(args, "target", "primary")

	cmd := fmt.Sprintf("find %s", shellQuote(path))
	if maxDepth > 0 {
		cmd += fmt.Sprintf(" -maxdepth %d", maxDepth)
	}
	if fileType == "f" || fileType == "d" {
		cmd += fmt.Sprintf(" -type %s", fileType)
	}
	cmd += fmt.Sprintf(" -name %s 2>/dev/null | head -100", shellQuote(pattern))

	output, err := m.Execute(ctx, cmd, target)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	return textResult(output), nil
}

func searchTextHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	m := manager()
	if m == nil {
		return errorResult("no SSH manager initialized"), nil
	}
	args := requestArgs(req)
	pattern, _ := requireString(args, "pattern")
	path, _ := requireString(args, "path")
	recursive := getBool(args, "recursive", false)
	ignoreCase := getBool(args, "ignore_case", false)
	ctxLines := getInt(args, "context", 0)
	target := getString(args, "target", "primary")

	cmd := "grep"
	if recursive {
		cmd += " -r"
	}
	if ignoreCase {
		cmd += " -i"
	}
	if ctxLines > 0 {
		cmd += fmt.Sprintf(" -C %d", ctxLines)
	}
	cmd += fmt.Sprintf(" -n %s %s 2>/dev/null | head -100", shellQuote(pattern), shellQuote(path))

	output, err := m.Execute(ctx, cmd, target)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	return textResult(output), nil
}

func packageManageHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	m := manager()
	if m == nil {
		return errorResult("no SSH manager initialized"), nil
	}
	args := requestArgs(req)
	action, _ := requireString(args, "action")
	pkg := getString(args, "package", "")
	target := getString(args, "target", "primary")

	if action != "list" && pkg == "" {
		return errorResult("Package name required for this action"), nil
	}

	detectCmd := `
if command -v apt-get >/dev/null 2>&1; then echo "apt"
elif command -v apk >/dev/null 2>&1; then echo "apk"
elif command -v dnf >/dev/null 2>&1; then echo "dnf"
elif command -v yum >/dev/null 2>&1; then echo "yum"
else echo "unknown"
fi`

	pkgMgr, err := m.Execute(ctx, detectCmd, target)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	pkgMgr = trimOutput(pkgMgr)

	var cmd string
	switch pkgMgr {
	case "apt":
		switch action {
		case "install":
			cmd = fmt.Sprintf("apt-get update && apt-get install -y %s", shellQuote(pkg))
		case "remove":
			cmd = fmt.Sprintf("apt-get remove -y %s", shellQuote(pkg))
		case "check":
			cmd = fmt.Sprintf("dpkg -s %s 2>/dev/null", shellQuote(pkg))
		case "list":
			cmd = "dpkg -l | head -50"
		}
	case "apk":
		switch action {
		case "install":
			cmd = fmt.Sprintf("apk add %s", shellQuote(pkg))
		case "remove":
			cmd = fmt.Sprintf("apk del %s", shellQuote(pkg))
		case "check":
			cmd = fmt.Sprintf("apk info %s 2>/dev/null", shellQuote(pkg))
		case "list":
			cmd = "apk list --installed | head -50"
		}
	case "dnf", "yum":
		switch action {
		case "install":
			cmd = fmt.Sprintf("%s install -y %s", pkgMgr, shellQuote(pkg))
		case "remove":
			cmd = fmt.Sprintf("%s remove -y %s", pkgMgr, shellQuote(pkg))
		case "check":
			cmd = fmt.Sprintf("rpm -qi %s 2>/dev/null", shellQuote(pkg))
		case "list":
			cmd = fmt.Sprintf("%s list installed 2>/dev/null | head -50", pkgMgr)
		}
	default:
		return errorResult("No supported package manager found"), nil
	}

	output, err := m.Execute(ctx, cmd, target)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	return textResult(output), nil
}