feat: add constrained read-only diagnostics
This commit is contained in:
@@ -98,6 +98,33 @@ func (b *MCPBackend) ExecCommand(ctx context.Context, clientIP, device, command,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (b *MCPBackend) DiagnosticCommand(ctx context.Context, clientIP, device string, req executor.DiagnosticRequest) (any, error) {
|
||||
if device == "" {
|
||||
device = b.server.LocalName
|
||||
}
|
||||
if err := b.authorize(device, "diagnostic", clientIP); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if req.Action == "" {
|
||||
return nil, errors.New("diagnostic action required")
|
||||
}
|
||||
if b.server.isRemote(device) {
|
||||
if b.server.Cluster == nil || !b.server.Cluster.Online(device) {
|
||||
return nil, fmt.Errorf("device %s is offline", device)
|
||||
}
|
||||
res, err := b.server.Cluster.Diagnostic(ctx, device, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
res, err := executor.Diagnostic(ctx, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (b *MCPBackend) ReadFile(ctx context.Context, clientIP, device, path string, offset, limit int64) (any, error) {
|
||||
if device == "" {
|
||||
device = b.server.LocalName
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/wjyhk/vps-commander/internal/executor"
|
||||
"github.com/wjyhk/vps-commander/internal/security"
|
||||
"github.com/wjyhk/vps-commander/internal/storage"
|
||||
"github.com/wjyhk/vps-commander/internal/sysinfo"
|
||||
@@ -464,6 +465,19 @@ func (m *Manager) Exec(ctx context.Context, name string, p ExecPayload) (Result,
|
||||
return Result{msg.ExitCode, msg.Stdout, msg.Stderr, msg.DurationMS}, nil
|
||||
}
|
||||
|
||||
func (m *Manager) Diagnostic(ctx context.Context, name string, req executor.DiagnosticRequest) (executor.DiagnosticResult, error) {
|
||||
msg, err := m.call(ctx, name, "diagnostic_command", req)
|
||||
if err != nil {
|
||||
return executor.DiagnosticResult{}, err
|
||||
}
|
||||
b, _ := json.Marshal(msg.Payload)
|
||||
var out executor.DiagnosticResult
|
||||
if err := json.Unmarshal(b, &out); err != nil {
|
||||
return out, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (m *Manager) ReadFile(ctx context.Context, name string, p FileReadPayload) (FileResult, error) {
|
||||
msg, err := m.call(ctx, name, "file_read", p)
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
package executor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// DiagnosticRequest is a structured, read-only diagnostic operation. It is deliberately
|
||||
// not a shell command: every action maps to a fixed executable/argument set.
|
||||
type DiagnosticRequest struct {
|
||||
Action string `json:"action"`
|
||||
Target string `json:"target,omitempty"`
|
||||
Lines int `json:"lines,omitempty"`
|
||||
}
|
||||
|
||||
type DiagnosticResult struct {
|
||||
Action string `json:"action"`
|
||||
Target string `json:"target,omitempty"`
|
||||
ExitCode int `json:"exit_code"`
|
||||
Stdout string `json:"stdout"`
|
||||
Stderr string `json:"stderr"`
|
||||
DurationMS int64 `json:"duration_ms"`
|
||||
}
|
||||
|
||||
func Diagnostic(ctx context.Context, req DiagnosticRequest) (DiagnosticResult, error) {
|
||||
args, err := diagnosticArgs(req)
|
||||
if err != nil {
|
||||
return DiagnosticResult{}, err
|
||||
}
|
||||
if req.Lines <= 0 {
|
||||
req.Lines = 100
|
||||
}
|
||||
if req.Lines > 500 {
|
||||
req.Lines = 500
|
||||
}
|
||||
// docker logs is the only operation where a caller-controlled numeric limit is useful.
|
||||
if req.Action == "docker_logs" {
|
||||
args = []string{"logs", "--tail", strconv.Itoa(req.Lines), req.Target}
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
c, cancel := context.WithTimeout(ctx, 15*time.Second)
|
||||
defer cancel()
|
||||
cmd := exec.CommandContext(c, args[0], args[1:]...)
|
||||
out, runErr := cmd.CombinedOutput()
|
||||
const maxOutput = 256 * 1024
|
||||
if len(out) > maxOutput {
|
||||
out = append(out[:maxOutput], []byte("\n[output truncated at 256 KiB]\n")...)
|
||||
}
|
||||
code := 0
|
||||
if cmd.ProcessState != nil {
|
||||
code = cmd.ProcessState.ExitCode()
|
||||
}
|
||||
result := DiagnosticResult{Action: req.Action, Target: req.Target, ExitCode: code, Stdout: string(out), DurationMS: time.Since(start).Milliseconds()}
|
||||
if runErr != nil && code == -1 {
|
||||
return result, fmt.Errorf("diagnostic failed: %w", runErr)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func diagnosticArgs(req DiagnosticRequest) ([]string, error) {
|
||||
target := strings.TrimSpace(req.Target)
|
||||
switch req.Action {
|
||||
case "docker_ps":
|
||||
if target != "" {
|
||||
return nil, errors.New("docker_ps does not accept target")
|
||||
}
|
||||
return []string{"docker", "ps", "--no-trunc"}, nil
|
||||
case "docker_logs":
|
||||
if target == "" || strings.ContainsAny(target, " \t\r\n;|&<>$`") {
|
||||
return nil, errors.New("invalid docker container target")
|
||||
}
|
||||
return []string{"docker", "logs", "--tail", "100", target}, nil
|
||||
case "docker_inspect":
|
||||
if target == "" || strings.ContainsAny(target, " \t\r\n;|&<>$`") {
|
||||
return nil, errors.New("invalid docker target")
|
||||
}
|
||||
return []string{"docker", "inspect", target}, nil
|
||||
case "systemctl_status":
|
||||
if target == "" || strings.ContainsAny(target, " \t\r\n;|&<>$`") {
|
||||
return nil, errors.New("invalid systemd unit")
|
||||
}
|
||||
return []string{"systemctl", "status", "--no-pager", "--full", target}, nil
|
||||
case "journalctl":
|
||||
if target == "" || strings.ContainsAny(target, " \t\r\n;|&<>$`") {
|
||||
return nil, errors.New("invalid journal unit")
|
||||
}
|
||||
return []string{"journalctl", "--no-pager", "-u", target, "-n", "100"}, nil
|
||||
case "ss_listen":
|
||||
if target != "" {
|
||||
return nil, errors.New("ss_listen does not accept target")
|
||||
}
|
||||
return []string{"ss", "-lntp"}, nil
|
||||
case "df":
|
||||
if target != "" {
|
||||
return nil, errors.New("df does not accept target")
|
||||
}
|
||||
return []string{"df", "-h"}, nil
|
||||
case "free":
|
||||
if target != "" {
|
||||
return nil, errors.New("free does not accept target")
|
||||
}
|
||||
return []string{"free", "-h"}, nil
|
||||
case "uptime":
|
||||
if target != "" {
|
||||
return nil, errors.New("uptime does not accept target")
|
||||
}
|
||||
return []string{"uptime"}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported diagnostic action: %s", req.Action)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package executor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDiagnosticRejectsShellInjection(t *testing.T) {
|
||||
_, err := Diagnostic(context.Background(), DiagnosticRequest{Action: "docker_logs", Target: "edgeever;rm -rf /"})
|
||||
if err == nil || !strings.Contains(err.Error(), "invalid docker") {
|
||||
t.Fatalf("expected target rejection, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiagnosticAllowsReadOnlyUptime(t *testing.T) {
|
||||
got, err := Diagnostic(context.Background(), DiagnosticRequest{Action: "uptime"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.ExitCode != 0 || strings.TrimSpace(got.Stdout) == "" {
|
||||
t.Fatalf("unexpected uptime result: %#v", got)
|
||||
}
|
||||
}
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/wjyhk/vps-commander/internal/auth"
|
||||
"github.com/wjyhk/vps-commander/internal/executor"
|
||||
)
|
||||
|
||||
// MCP Protocol / JSON-RPC 2.0 Types
|
||||
@@ -77,6 +78,10 @@ type ActionBackend interface {
|
||||
EditBlock(ctx context.Context, clientIP, device, path, oldText, newText string) (any, error)
|
||||
}
|
||||
|
||||
type DiagnosticBackend interface {
|
||||
DiagnosticCommand(ctx context.Context, clientIP, device string, req executor.DiagnosticRequest) (any, error)
|
||||
}
|
||||
|
||||
type ExtendedBackend interface {
|
||||
ReadMultipleFiles(ctx context.Context, clientIP, device string, paths []string, limit int64) (any, error)
|
||||
CreateDirectory(ctx context.Context, clientIP, device, path string) (any, error)
|
||||
@@ -280,6 +285,7 @@ func (s *Server) handleRequest(ctx context.Context, r *http.Request, req JSONRPC
|
||||
{Name: "create_directory", Description: "创建目录", InputSchema: InputSchema{Type: "object", Properties: map[string]PropertyDef{"device": {Type: "string", Description: "目标设备"}, "path": {Type: "string", Description: "目录路径"}}, Required: []string{"path"}}},
|
||||
{Name: "move_file", Description: "移动或重命名文件", InputSchema: InputSchema{Type: "object", Properties: map[string]PropertyDef{"device": {Type: "string", Description: "目标设备"}, "source": {Type: "string", Description: "源路径"}, "destination": {Type: "string", Description: "目标路径"}}, Required: []string{"source", "destination"}}},
|
||||
{Name: "list_processes", Description: "列出受控节点当前进程", InputSchema: InputSchema{Type: "object", Properties: map[string]PropertyDef{"device": {Type: "string", Description: "目标设备"}}}},
|
||||
{Name: "diagnostic_command", Description: "执行严格白名单的只读系统诊断,不接受任意 Shell", InputSchema: InputSchema{Type: "object", Properties: map[string]PropertyDef{"device": {Type: "string", Description: "目标设备"}, "action": {Type: "string", Description: "docker_ps|docker_logs|docker_inspect|systemctl_status|journalctl|ss_listen|df|free|uptime"}, "target": {Type: "string", Description: "容器名或 systemd unit(按 action 使用)"}, "lines": {Type: "integer", Description: "docker_logs 最大行数,1-500"}}, Required: []string{"action"}}},
|
||||
{
|
||||
Name: "edit_block",
|
||||
Description: "对指定文件执行唯一文本块的原子替换;old_text 必须恰好匹配一次",
|
||||
@@ -430,6 +436,24 @@ func (s *Server) callTool(ctx context.Context, clientIP string, params CallToolP
|
||||
return map[string]string{"error": err.Error()}, true
|
||||
}
|
||||
return res, false
|
||||
case "diagnostic_command":
|
||||
var a struct {
|
||||
Device string `json:"device"`
|
||||
Action string `json:"action"`
|
||||
Target string `json:"target"`
|
||||
Lines int `json:"lines"`
|
||||
}
|
||||
_ = json.Unmarshal(params.Arguments, &a)
|
||||
b, ok := s.Backend.(DiagnosticBackend)
|
||||
if !ok {
|
||||
return map[string]string{"error": "diagnostic tools unavailable"}, true
|
||||
}
|
||||
res, err := b.DiagnosticCommand(ctx, clientIP, a.Device, executor.DiagnosticRequest{Action: a.Action, Target: a.Target, Lines: a.Lines})
|
||||
if err != nil {
|
||||
return map[string]string{"error": err.Error()}, true
|
||||
}
|
||||
return res, false
|
||||
|
||||
case "list_processes":
|
||||
var a struct {
|
||||
Device string `json:"device"`
|
||||
|
||||
@@ -79,7 +79,7 @@ func TestMCPInitializeAndListTools(t *testing.T) {
|
||||
t.Fatalf("expected map result")
|
||||
}
|
||||
tools, ok := resMap["tools"].([]any)
|
||||
if !ok || len(tools) != 11 {
|
||||
if !ok || len(tools) != 12 {
|
||||
t.Fatalf("expected 11 tools, got %d", len(tools))
|
||||
}
|
||||
foundWriteFile := false
|
||||
|
||||
@@ -15,6 +15,7 @@ var registry = map[string]ToolSpec{
|
||||
"start_search": {Name: "start_search", RequiredRisk: RiskLow, Category: "search", ReadOnly: true},
|
||||
"get_more_search_results": {Name: "get_more_search_results", RequiredRisk: RiskLow, Category: "search", ReadOnly: true},
|
||||
"list_processes": {Name: "list_processes", RequiredRisk: RiskLow, Category: "process", ReadOnly: true},
|
||||
"diagnostic_command": {Name: "diagnostic_command", RequiredRisk: RiskLow, Category: "diagnostic", ReadOnly: true},
|
||||
"list_sessions": {Name: "list_sessions", RequiredRisk: RiskLow, Category: "process", ReadOnly: true},
|
||||
"read_process_output": {Name: "read_process_output", RequiredRisk: RiskLow, Category: "process", ReadOnly: true},
|
||||
"edit_block": {Name: "edit_block", RequiredRisk: RiskMedium, Category: "file", ReadOnly: false},
|
||||
|
||||
Reference in New Issue
Block a user