model.go

package main

import (
	"fmt"
	"math"
	"math/rand"
	"net/netip"
	"sort"
	"time"
)

type AR1 struct {
	Correlation float64
	Sigma       float64
	Value       float64
}

func (a *AR1) Step(rng *rand.Rand) float64 {
	if a.Sigma == 0 {
		a.Value = 0
		return 0
	}
	innovationScale := a.Sigma * math.Sqrt(1-a.Correlation*a.Correlation)
	a.Value = a.Correlation*a.Value + innovationScale*rng.NormFloat64()
	return a.Value
}

type SampleBudget struct {
	Bytes         float64
	NextFrameSize int
}

func (b *SampleBudget) AddRate(rateBPS float64, interval time.Duration, samplingRate uint32) {
	b.Bytes += rateBPS / 8 * interval.Seconds() / float64(samplingRate)
}

func (b *SampleBudget) Take(sampler *PacketSizeSampler, minimum int) (int, bool) {
	if b.NextFrameSize == 0 || b.NextFrameSize < minimum {
		b.NextFrameSize = sampler.Next(minimum)
	}
	if b.Bytes < float64(b.NextFrameSize) {
		return 0, false
	}
	frameSize := b.NextFrameSize
	b.Bytes -= float64(frameSize)
	b.NextFrameSize = sampler.Next(minimum)
	return frameSize, true
}

type AmbientHostState struct {
	Address          netip.Addr
	BaseWeight       float64
	Noise            AR1
	IncomingBudget   SampleBudget
	OutgoingBudget   SampleBudget
	LastRate         float64
	LastIncomingRate float64
	LastOutgoingRate float64
}

type AmbientNetworkState struct {
	Config      RuntimeAmbientNetwork
	GlobalNoise AR1
	Hosts       []AmbientHostState
}

type AmbientModel struct {
	Networks     []AmbientNetworkState
	RNG          *rand.Rand
	Sampler      PacketSizeSampler
	Chooser      ProtocolChooser
	SamplingRate uint32
}

type PacketSpec struct {
	Internal        netip.Addr
	Peer            netip.Addr
	Direction       Direction
	Protocol        Protocol
	FrameLength     int
	DestinationPort uint16
	TCPFlags        string
}

type TickResult struct {
	Packets           []PacketSpec
	TargetBPS         float64
	TargetIncomingBPS float64
	TargetOutgoingBPS float64
	TopAddress        netip.Addr
	TopRateBPS        float64
}

func NewAmbientModel(cfg RuntimeAmbient) (*AmbientModel, error) {
	rng := rand.New(rand.NewSource(cfg.Common.Seed))
	chooser, err := NewProtocolChooser(cfg.Common.Protocols)
	if err != nil {
		return nil, err
	}
	model := &AmbientModel{
		RNG:          rng,
		Sampler:      NewPacketSizeSampler(cfg.Common.Packet, rng),
		Chooser:      chooser,
		SamplingRate: cfg.Common.SamplingRate,
	}

	for _, networkCfg := range cfg.Networks {
		addresses, err := generateAddresses(networkCfg.Prefix, networkCfg.Hosts, rng)
		if err != nil {
			return nil, err
		}
		state := AmbientNetworkState{
			Config: networkCfg,
			GlobalNoise: AR1{
				Correlation: networkCfg.Correlation,
				Sigma:       networkCfg.TotalNoise,
			},
		}
		for rank, address := range addresses {
			weight := 1.0
			if networkCfg.Zipf > 0 {
				weight = 1 / math.Pow(float64(rank+1), networkCfg.Zipf)
			}
			state.Hosts = append(state.Hosts, AmbientHostState{
				Address:    address,
				BaseWeight: weight,
				Noise: AR1{
					Correlation: networkCfg.Correlation,
					Sigma:       networkCfg.HostNoise,
				},
			})
		}
		model.Networks = append(model.Networks, state)
	}
	return model, nil
}

func (m *AmbientModel) Tick(interval time.Duration, maxSamples int) (TickResult, error) {
	result := TickResult{Packets: make([]PacketSpec, 0, 128)}

	for networkIndex := range m.Networks {
		network := &m.Networks[networkIndex]
		globalFactor := math.Exp(network.GlobalNoise.Step(m.RNG))
		globalFactor = clamp(globalFactor, 0.80, 1.20)
		networkRate := network.Config.RateBPS * globalFactor
		networkIncomingRate := networkRate * network.Config.IncomingShare
		networkOutgoingRate := networkRate * network.Config.OutgoingShare
		result.TargetBPS += networkRate
		result.TargetIncomingBPS += networkIncomingRate
		result.TargetOutgoingBPS += networkOutgoingRate

		scores := make([]float64, len(network.Hosts))
		scoreTotal := 0.0
		for i := range network.Hosts {
			host := &network.Hosts[i]
			score := host.BaseWeight * math.Exp(host.Noise.Step(m.RNG))
			scores[i] = score
			scoreTotal += score
		}

		for i := range network.Hosts {
			host := &network.Hosts[i]
			hostShare := scores[i] / scoreTotal
			host.LastIncomingRate = networkIncomingRate * hostShare
			host.LastOutgoingRate = networkOutgoingRate * hostShare
			host.LastRate = host.LastIncomingRate + host.LastOutgoingRate
			if host.LastRate > result.TopRateBPS {
				result.TopRateBPS = host.LastRate
				result.TopAddress = host.Address
			}

			host.IncomingBudget.AddRate(host.LastIncomingRate, interval, m.SamplingRate)
			host.OutgoingBudget.AddRate(host.LastOutgoingRate, interval, m.SamplingRate)

			for len(result.Packets) < maxSamples {
				generated := false
				for _, direction := range []Direction{DirectionIncoming, DirectionOutgoing} {
					if len(result.Packets) >= maxSamples {
						break
					}
					packet, ok, err := m.takeAmbientPacket(host, network, direction)
					if err != nil {
						return TickResult{}, err
					}
					if ok {
						result.Packets = append(result.Packets, packet)
						generated = true
					}
				}
				if !generated {
					break
				}
			}
			if len(result.Packets) >= maxSamples {
				return result, nil
			}
		}
	}
	return result, nil
}

func (m *AmbientModel) takeAmbientPacket(host *AmbientHostState, network *AmbientNetworkState, direction Direction) (PacketSpec, bool, error) {
	budget := &host.IncomingBudget
	if direction == DirectionOutgoing {
		budget = &host.OutgoingBudget
	}

	protocol := m.Chooser.Next(m.RNG)
	minimum := minimumFrameLength(host.Address.Is6(), protocol)
	frameSize, ok := budget.Take(&m.Sampler, minimum)
	if !ok {
		return PacketSpec{}, false, nil
	}
	peer, err := randomPrefixAddress(network.Config.PeerPrefixes, m.RNG)
	if err != nil {
		return PacketSpec{}, false, err
	}
	return PacketSpec{
		Internal:    host.Address,
		Peer:        peer,
		Direction:   direction,
		Protocol:    protocol,
		FrameLength: frameSize,
		TCPFlags:    "ack",
	}, true, nil
}

type AttackModel struct {
	Config  RuntimeAttack
	RNG     *rand.Rand
	Sampler PacketSizeSampler
	Chooser ProtocolChooser
	Budget  SampleBudget
}

func NewAttackModel(cfg RuntimeAttack) (*AttackModel, error) {
	chooser, err := NewProtocolChooser(cfg.Common.Protocols)
	if err != nil {
		return nil, err
	}
	rng := rand.New(rand.NewSource(cfg.Common.Seed))
	return &AttackModel{
		Config:  cfg,
		RNG:     rng,
		Sampler: NewPacketSizeSampler(cfg.Common.Packet, rng),
		Chooser: chooser,
	}, nil
}

func (m *AttackModel) TotalDuration() time.Duration {
	return m.Config.RampUp + m.Config.Hold + m.Config.RampDown
}

func (m *AttackModel) RateAt(elapsed time.Duration) float64 {
	if elapsed < 0 {
		return 0
	}
	if m.Config.RampUp > 0 && elapsed < m.Config.RampUp {
		return m.Config.PeakRateBPS * smoothstep(float64(elapsed)/float64(m.Config.RampUp))
	}
	elapsed -= m.Config.RampUp
	if elapsed < m.Config.Hold {
		return m.Config.PeakRateBPS
	}
	elapsed -= m.Config.Hold
	if m.Config.RampDown > 0 && elapsed < m.Config.RampDown {
		return m.Config.PeakRateBPS * (1 - smoothstep(float64(elapsed)/float64(m.Config.RampDown)))
	}
	return 0
}

func (m *AttackModel) Tick(elapsed, interval time.Duration, maxSamples int) (TickResult, error) {
	rate := m.RateAt(elapsed)
	result := TickResult{TargetBPS: rate, TopAddress: m.Config.Victim, TopRateBPS: rate}
	m.Budget.AddRate(rate, interval, m.Config.Common.SamplingRate)

	for len(result.Packets) < maxSamples {
		protocol := m.Chooser.Next(m.RNG)
		minimum := minimumFrameLength(m.Config.Victim.Is6(), protocol)
		frameSize, ok := m.Budget.Take(&m.Sampler, minimum)
		if !ok {
			break
		}
		peer, err := randomPrefixAddress(m.Config.SourcePrefixes, m.RNG)
		if err != nil {
			return TickResult{}, err
		}
		result.Packets = append(result.Packets, PacketSpec{
			Internal:        m.Config.Victim,
			Peer:            peer,
			Direction:       m.Config.Direction,
			Protocol:        protocol,
			FrameLength:     frameSize,
			DestinationPort: m.Config.DestinationPort,
			TCPFlags:        m.Config.TCPFlags,
		})
	}
	return result, nil
}

func randomPrefixAddress(prefixes []netip.Prefix, rng *rand.Rand) (netip.Addr, error) {
	if len(prefixes) == 0 {
		return netip.Addr{}, fmt.Errorf("no address prefixes configured")
	}
	return randomAddress(prefixes[rng.Intn(len(prefixes))], rng)
}

func smoothstep(value float64) float64 {
	value = clamp(value, 0, 1)
	return value * value * (3 - 2*value)
}

func clamp(value, minimum, maximum float64) float64 {
	return math.Max(minimum, math.Min(maximum, value))
}

func sortedTopHosts(networks []AmbientNetworkState, count int) []AmbientHostState {
	var all []AmbientHostState
	for _, network := range networks {
		all = append(all, network.Hosts...)
	}
	sort.Slice(all, func(i, j int) bool { return all[i].LastRate > all[j].LastRate })
	if len(all) > count {
		all = all[:count]
	}
	return all
}