ProxmoxDash/traffic_test.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func trafficFixture(t *testing.T) (*Store, agentConfig) {
	t.Helper()
	s, err := openStore(":memory:")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { s.Close() })
	token, err := s.CreateAgentEnrollmentToken()
	if err != nil {
		t.Fatal(err)
	}
	a, err := s.EnrollAgent(agentEnrollmentRequest{Token: token, Name: "NPM test", VMID: 102})
	if err != nil {
		t.Fatal(err)
	}
	return s, a
}
func npmTestLine(ip, path string) string {
	return fmt.Sprintf("[%s] - - 301 - GET https example.test \"%s\" [Client %s] [Length 166] [Gzip -] [Sent-to 192.168.0.9] \"-\" \"-\"\n", time.Now().UTC().Format("02/Jan/2006:15:04:05 -0700"), path, ip)
}
func TestNPMTrafficFormats(t *testing.T) {
	for _, status := range []string{"- 200 200", "- - 301", "- 502 502", "- - 401"} {
		line := strings.Replace(npmTestLine("2001:4860:4860::8888", "/login?token=secret#fragment"), "- - 301", status, 1)
		e, ok := parseNPMTraffic(line)
		if !ok || e.Kind != "request" || e.Path != "/login" || e.Domain != "example.test" || e.IP != "2001:4860:4860::8888" {
			t.Fatalf("bad parse: %+v %v", e, ok)
		}
	}
	if _, ok := parseNPMTraffic("broken"); ok {
		t.Fatal("accepted invalid log")
	}
}
func TestCrowdTrafficDoesNotInventBlockedRequests(t *testing.T) {
	raw := `[{"id":12,"uuid":"alert-uuid","created_at":"2026-10-09T16:16:43Z","source":{"ip":"20.194.96.114"},"scenario":"crowdsecurity/http-admin-interface-probing","events":[{"meta":[{"key":"target_fqdn","value":"example.test"}]}],"decisions":[{"id":42,"value":"20.194.96.114","type":"ban","simulated":false},{"id":43,"value":"1.1.1.1","type":"ban","simulated":true}]}]`
	events, err := parseCrowdTraffic([]byte(raw))
	if err != nil || len(events) != 2 {
		t.Fatalf("%+v %v", events, err)
	}
	if events[0].Kind != "alert" || events[1].Kind != "decision" || events[1].Domain != "example.test" {
		t.Fatalf("bad categories %+v", events)
	}
	simulated := strings.Replace(raw, `"id":12`, `"simulated":true,"id":12`, 1)
	events, err = parseCrowdTraffic([]byte(simulated))
	if err != nil || len(events) != 0 {
		t.Fatal("simulation counted as live")
	}
}
func TestTrafficLogCheckpointRotationAndPartialLines(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "proxy-host-1_access.log")
	os.WriteFile(path, []byte(npmTestLine("1.1.1.1", "/old")), 0600)
	state := trafficAgentState{Cursors: map[string]trafficCursor{}}
	if _, err := state.readLogs(dir); err != nil {
		t.Fatal(err)
	}
	if len(state.Pending) != 0 {
		t.Fatal("imported old logs on first activation")
	}
	appendFile := func(path, s string) {
		f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0600)
		if err != nil {
			t.Fatal(err)
		}
		f.WriteString(s)
		f.Close()
	}
	line := npmTestLine("8.8.8.8", "/new")
	appendFile(path, line[:len(line)/2])
	state.readLogs(dir)
	if len(state.Pending) != 0 {
		t.Fatal("partial line emitted")
	}
	appendFile(path, line[len(line)/2:])
	state.readLogs(dir)
	if len(state.Pending) != 1 {
		t.Fatal("complete line missing")
	}
	id := state.Pending[0].ID
	state.Pending = nil
	checkpoint := filepath.Join(dir, "state.json")
	if err := saveTrafficAgentState(checkpoint, state); err != nil {
		t.Fatal(err)
	}
	raw, _ := os.ReadFile(checkpoint)
	var restored trafficAgentState
	json.Unmarshal(raw, &restored)
	os.Rename(path, path+".1")
	appendFile(path+".1", npmTestLine("9.9.9.9", "/rotated"))
	os.WriteFile(path, []byte(npmTestLine("1.0.0.1", "/fresh")), 0600)
	restored.readLogs(dir)
	if len(restored.Pending) != 2 {
		t.Fatalf("rotation lost events %+v", restored)
	}
	for _, e := range restored.Pending {
		if e.ID == id {
			t.Fatal("duplicate id")
		}
	}
	restored.Pending = nil
	restored.readLogs(dir)
	if len(restored.Pending) != 0 {
		t.Fatal("re-read same lines")
	}
}
func TestTrafficStoreDedupFiltersAndSanitization(t *testing.T) {
	s, a := trafficFixture(t)
	if s.TrafficSettings().Enabled {
		t.Fatal("enabled by default")
	}
	if _, err := s.SaveTrafficSettings(TrafficSettings{Enabled: true, AgentID: a.ID}); err != nil {
		t.Fatal(err)
	}
	now := time.Now().Unix()
	batch := TrafficBatch{Files: 22, Events: []TrafficEvent{{ID: "1", Time: now, Kind: "request", IP: "1.1.1.1", Domain: "example.test", Path: "/login?secret=x", Status: 502}, {ID: "2", Time: now, Kind: "request", IP: "192.168.0.10", Domain: "example.test", Status: 200}, {ID: "3", Time: now, Kind: "alert", IP: "1.1.1.1", Reason: "probe"}, {ID: "4", Time: now, Kind: "decision", IP: "1.1.1.1", Reason: "ban"}}}
	for i := 0; i < 2; i++ {
		if err := s.SaveTraffic(a.ID, batch); err != nil {
			t.Fatal(err)
		}
	}
	v, err := s.TrafficView("", "", false)
	if err != nil {
		t.Fatal(err)
	}
	if v.Requests != 2 || v.UniqueIPs != 2 || v.Alerts != 1 || v.Decisions != 1 || len(v.Events) != 4 {
		t.Fatalf("wrong counts: %+v", v)
	}
	for _, e := range v.Events {
		if strings.Contains(e.Path, "secret") {
			t.Fatal("secret saved")
		}
	}
	v, err = s.TrafficView("example.test", "request", true)
	if err != nil || v.Requests != 1 {
		t.Fatalf("bad filter %+v %v", v, err)
	}
	batch.Events = []TrafficEvent{{ID: "invalid", Kind: "blocked", IP: "1.1.1.1", Time: now}}
	if s.SaveTraffic(a.ID, batch) == nil {
		t.Fatal("accepted invented block event")
	}
}
func TestTrafficAPIOptInAndSourceAuthorization(t *testing.T) {
	s, a := trafficFixture(t)
	mux := http.NewServeMux()
	trafficRoutes(mux, s)
	get := func(path string) int {
		r := httptest.NewRecorder()
		mux.ServeHTTP(r, httptest.NewRequest("GET", path, nil))
		return r.Code
	}
	if get("/traffic") != 404 || get("/api/traffic") != 404 {
		t.Fatal("disabled page exposed")
	}
	send := func(secret string) int {
		r := httptest.NewRecorder()
		req := httptest.NewRequest("POST", "/api/agent/traffic", strings.NewReader(`{"events":[]}`))
		req.Header.Set("Authorization", "Bearer "+a.ID+"."+secret)
		mux.ServeHTTP(r, req)
		return r.Code
	}
	if send("wrong") != 401 || send(a.Secret) != 403 {
		t.Fatal("disabled/unauthorized ingestion accepted")
	}
	if _, err := s.SaveTrafficSettings(TrafficSettings{Enabled: true, AgentID: "missing"}); err == nil {
		t.Fatal("unknown agent accepted")
	}
	s.SaveTrafficSettings(TrafficSettings{Enabled: true, AgentID: a.ID, ServerCountry: "RU"})
	if send(a.Secret) != 200 || get("/traffic") != 200 {
		t.Fatal("selected agent not working")
	}
	token, _ := s.CreateAgentEnrollmentToken()
	other, _ := s.EnrollAgent(agentEnrollmentRequest{Token: token, Name: "Other"})
	s.SaveTrafficSettings(TrafficSettings{Enabled: true, AgentID: other.ID})
	if send(a.Secret) != 403 {
		t.Fatal("old source still accepted")
	}
	s.SaveTrafficSettings(TrafficSettings{Enabled: false, AgentID: other.ID})
	if get("/api/traffic") != 404 {
		t.Fatal("disable failed")
	}
}
func TestTrafficGeoDatabase(t *testing.T) {
	if !trafficGeoReady() {
		t.Fatal("embedded database invalid")
	}
	if lookupTrafficCountry("192.168.0.10") != "LAN" {
		t.Fatal("LAN mapped to country")
	}
	if c := lookupTrafficCountry("8.8.8.8"); len(c) != 2 {
		t.Fatalf("bad public IP geo: %q", c)
	}
}
func TestTrafficCheckpointKeepsPendingUntilAcknowledged(t *testing.T) {
	path := filepath.Join(t.TempDir(), "state.json")
	s := trafficAgentState{Pending: []TrafficEvent{{ID: "retry-me", IP: "1.1.1.1"}}, Cursors: map[string]trafficCursor{"file": {Offset: 100}}}
	if err := saveTrafficAgentState(path, s); err != nil {
		t.Fatal(err)
	}
	raw, _ := os.ReadFile(path)
	if !bytes.Contains(raw, []byte("retry-me")) {
		t.Fatal("pending batch lost")
	}
}