first commit

traffic/tracker.go

@@ -0,0 +1,194 @@
+package traffic
+
+import (
+	"bufio"
+	"fmt"
+	"os/exec"
+	"regexp"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+)
+
+const (
+	pollInterval     = 20 * time.Second
+	OnlineWindow     = 5 * time.Minute
+	trackerTableName = "alpine-router-traffic"
+)
+
+type IPStats struct {
+	TxBytes    uint64
+	RxBytes    uint64
+	LastActive time.Time
+}
+
+var (
+	mu     sync.RWMutex
+	stats  = map[string]*IPStats{}
+	prev   = map[string][2]uint64{}
+	useNFT bool
+)
+
+func Available() bool {
+	mu.RLock()
+	v := useNFT
+	mu.RUnlock()
+	return v
+}
+
+func Start() {
+	if _, err := exec.LookPath("nft"); err != nil {
+		return
+	}
+	if err := setupNFTTable(); err != nil {
+		return
+	}
+	mu.Lock()
+	useNFT = true
+	mu.Unlock()
+
+	go func() {
+		poll()
+		t := time.NewTicker(pollInterval)
+		defer t.Stop()
+		for range t.C {
+			poll()
+		}
+	}()
+}
+
+func Get(ip string) IPStats {
+	mu.RLock()
+	defer mu.RUnlock()
+	if s, ok := stats[ip]; ok {
+		return *s
+	}
+	return IPStats{}
+}
+
+func IsOnline(ip string) bool {
+	s := Get(ip)
+	return !s.LastActive.IsZero() && time.Since(s.LastActive) < OnlineWindow
+}
+
+func EnsureIPTracked(ip string) {
+	if ip == "" {
+		return
+	}
+	mu.Lock()
+	defer mu.Unlock()
+	if _, ok := prev[ip]; ok {
+		return
+	}
+	prev[ip] = [2]uint64{}
+	addNFTRule(ip)
+}
+
+func setupNFTTable() error {
+	exec.Command("nft", "delete", "table", "ip", trackerTableName).Run()
+
+	script := fmt.Sprintf(`table ip %s {
+	chain tx {
+		type filter hook forward priority filter + 10; policy accept;
+	}
+	chain rx {
+		type filter hook forward priority filter + 20; policy accept;
+	}
+}`, trackerTableName)
+
+	cmd := exec.Command("nft", "-f", "-")
+	cmd.Stdin = strings.NewReader(script)
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		return fmt.Errorf("nft setup: %s: %w", strings.TrimSpace(string(out)), err)
+	}
+	return nil
+}
+
+func addNFTRule(ip string) {
+	exec.Command("nft", "add", "rule", "ip", trackerTableName, "tx",
+		"ip", "saddr", ip, "counter").Run()
+	exec.Command("nft", "add", "rule", "ip", trackerTableName, "rx",
+		"ip", "daddr", ip, "counter").Run()
+}
+
+var txCounterRe = regexp.MustCompile(`ip saddr (\S+) counter packets \d+ bytes (\d+)`)
+var rxCounterRe = regexp.MustCompile(`ip daddr (\S+) counter packets \d+ bytes (\d+)`)
+
+func poll() {
+	mu.RLock()
+	nft := useNFT
+	mu.RUnlock()
+	if !nft {
+		return
+	}
+
+	current, err := readNFTCounters()
+	if err != nil {
+		return
+	}
+
+	now := time.Now()
+	mu.Lock()
+	defer mu.Unlock()
+
+	for ip, cur := range current {
+		p := prev[ip]
+		dTx := delta(p[0], cur[0])
+		dRx := delta(p[1], cur[1])
+
+		s := stats[ip]
+		if s == nil {
+			s = &IPStats{}
+			stats[ip] = s
+		}
+		s.TxBytes += dTx
+		s.RxBytes += dRx
+		if dTx > 0 || dRx > 0 {
+			s.LastActive = now
+		}
+		prev[ip] = cur
+	}
+
+	for ip := range prev {
+		if _, exists := current[ip]; !exists {
+			prev[ip] = [2]uint64{0, 0}
+		}
+	}
+}
+
+func readNFTCounters() (map[string][2]uint64, error) {
+	cmd := exec.Command("nft", "list", "table", "ip", trackerTableName)
+	out, err := cmd.Output()
+	if err != nil {
+		return nil, err
+	}
+
+	result := map[string][2]uint64{}
+	scanner := bufio.NewScanner(strings.NewReader(string(out)))
+	for scanner.Scan() {
+		line := scanner.Text()
+		if m := txCounterRe.FindStringSubmatch(line); m != nil {
+			ip := m[1]
+			b, _ := strconv.ParseUint(m[2], 10, 64)
+			cur := result[ip]
+			cur[0] += b
+			result[ip] = cur
+		} else if m := rxCounterRe.FindStringSubmatch(line); m != nil {
+			ip := m[1]
+			b, _ := strconv.ParseUint(m[2], 10, 64)
+			cur := result[ip]
+			cur[1] += b
+			result[ip] = cur
+		}
+	}
+	return result, nil
+}
+
+func delta(old, cur uint64) uint64 {
+	if cur >= old {
+		return cur - old
+	}
+	return cur
+}