runner.go

package main

import (
	"context"
	"fmt"
	"log"
	"time"
)

func runAmbient(parent context.Context, cfg RuntimeAmbient) error {
	ctx := parent
	var cancel context.CancelFunc
	if cfg.Duration > 0 {
		ctx, cancel = context.WithTimeout(parent, cfg.Duration)
		defer cancel()
	}

	model, err := NewAmbientModel(cfg)
	if err != nil {
		return err
	}
	exporter, err := NewSFlowExporter(cfg.Common)
	if err != nil {
		return err
	}
	defer func() {
		if err := exporter.Close(); err != nil {
			log.Printf("close exporter: %v", err)
		}
	}()

	log.Printf("mode=ambient collector=%s agent=%s sub_agent=%d sampling=%d networks=%d duration=%s",
		cfg.Common.Collector, cfg.Common.AgentIP, cfg.Common.SubAgentID,
		cfg.Common.SamplingRate, len(cfg.Networks), printableDuration(cfg.Duration))

	ticker := time.NewTicker(cfg.Common.Tick)
	defer ticker.Stop()
	lastLog := time.Now()

	for {
		select {
		case <-ctx.Done():
			if parent.Err() != nil {
				return parent.Err()
			}
			return nil
		case <-ticker.C:
			result, err := model.Tick(cfg.Common.Tick, cfg.Common.MaxSamplesPerTick)
			if err != nil {
				return err
			}
			for _, packet := range result.Packets {
				if err := exporter.Add(packet); err != nil {
					return err
				}
			}
			if err := exporter.Flush(); err != nil {
				return err
			}
			if len(result.Packets) >= cfg.Common.MaxSamplesPerTick {
				log.Printf("warning: ambient reached max_samples_per_tick=%d; generated rate may be clipped", cfg.Common.MaxSamplesPerTick)
			}

			if time.Since(lastLog) >= cfg.Common.LogInterval {
				interval := time.Since(lastLog)
				stats := exporter.SnapshotAndReset()
				top := sortedTopHosts(model.Networks, 1)
				topText := "none"
				if len(top) > 0 {
					topText = fmt.Sprintf("%s@%s", top[0].Address, formatRate(top[0].LastRate))
				}
				log.Printf("mode=ambient target=%s target_in=%s target_out=%s estimated=%s estimated_in=%s estimated_out=%s samples=%d datagrams=%d top=%s",
					formatRate(result.TargetBPS), formatRate(result.TargetIncomingBPS), formatRate(result.TargetOutgoingBPS),
					measuredRate(stats.EstimatedBits, interval), measuredRate(stats.EstimatedIncomingBits, interval), measuredRate(stats.EstimatedOutgoingBits, interval),
					stats.Samples, stats.Datagrams, topText)
				lastLog = time.Now()
			}
		}
	}
}

func runAttack(ctx context.Context, cfg RuntimeAttack) error {
	model, err := NewAttackModel(cfg)
	if err != nil {
		return err
	}
	exporter, err := NewSFlowExporter(cfg.Common)
	if err != nil {
		return err
	}
	defer func() {
		if err := exporter.Close(); err != nil {
			log.Printf("close exporter: %v", err)
		}
	}()

	totalDuration := model.TotalDuration()
	log.Printf("mode=attack collector=%s agent=%s sub_agent=%d sampling=%d victim=%s peak=%s duration=%s",
		cfg.Common.Collector, cfg.Common.AgentIP, cfg.Common.SubAgentID,
		cfg.Common.SamplingRate, cfg.Victim, formatRate(cfg.PeakRateBPS), totalDuration)

	start := time.Now()
	lastTick := start
	lastLog := start
	ticker := time.NewTicker(cfg.Common.Tick)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case now := <-ticker.C:
			elapsed := now.Sub(start)
			interval := now.Sub(lastTick)
			lastTick = now
			if elapsed > totalDuration {
				if err := exporter.Flush(); err != nil {
					return err
				}
				stats := exporter.SnapshotAndReset()
				if stats.Samples > 0 {
					log.Printf("mode=attack target=0bps estimated=%s samples=%d datagrams=%d phase=done",
						measuredRate(stats.EstimatedBits, now.Sub(lastLog)), stats.Samples, stats.Datagrams)
				}
				return nil
			}

			result, err := model.Tick(elapsed, interval, cfg.Common.MaxSamplesPerTick)
			if err != nil {
				return err
			}
			for _, packet := range result.Packets {
				if err := exporter.Add(packet); err != nil {
					return err
				}
			}
			if err := exporter.Flush(); err != nil {
				return err
			}
			if len(result.Packets) >= cfg.Common.MaxSamplesPerTick {
				log.Printf("warning: attack reached max_samples_per_tick=%d; generated rate may be clipped", cfg.Common.MaxSamplesPerTick)
			}

			if now.Sub(lastLog) >= cfg.Common.LogInterval {
				stats := exporter.SnapshotAndReset()
				log.Printf("mode=attack victim=%s target=%s estimated=%s samples=%d datagrams=%d phase=%s",
					cfg.Victim, formatRate(result.TargetBPS), measuredRate(stats.EstimatedBits, now.Sub(lastLog)),
					stats.Samples, stats.Datagrams, attackPhase(elapsed, cfg))
				lastLog = now
			}
		}
	}
}

func attackPhase(elapsed time.Duration, cfg RuntimeAttack) string {
	if elapsed < cfg.RampUp {
		return "ramp-up"
	}
	elapsed -= cfg.RampUp
	if elapsed < cfg.Hold {
		return "hold"
	}
	return "ramp-down"
}

func measuredRate(bits uint64, interval time.Duration) string {
	if interval <= 0 {
		return "0bps"
	}
	return formatRate(float64(bits) / interval.Seconds())
}

func formatRate(bitsPerSecond float64) string {
	switch {
	case bitsPerSecond >= 1e12:
		return fmt.Sprintf("%.2fTbps", bitsPerSecond/1e12)
	case bitsPerSecond >= 1e9:
		return fmt.Sprintf("%.2fGbps", bitsPerSecond/1e9)
	case bitsPerSecond >= 1e6:
		return fmt.Sprintf("%.2fMbps", bitsPerSecond/1e6)
	case bitsPerSecond >= 1e3:
		return fmt.Sprintf("%.2fKbps", bitsPerSecond/1e3)
	default:
		return fmt.Sprintf("%.0fbps", bitsPerSecond)
	}
}

func printableDuration(duration time.Duration) string {
	if duration == 0 {
		return "infinite"
	}
	return duration.String()
}