tamabox/internal/security/autoban/autoban_test.go

package autoban

import (
	"sync/atomic"
	"testing"
	"time"
)

func newTestBanner(thresholdSensitive, threshold404 int, window, duration time.Duration, alertEmail string, alertCalls *int64) *Banner {
	return New(Options{
		Enabled:            true,
		TrackNotFound:      true, // 测试 helper 显式开启 404 跟踪，覆盖两种计数的既有用例
		ThresholdSensitive: thresholdSensitive,
		Threshold404:       threshold404,
		Window:             window,
		Duration:           duration,
		AlertEmail:         alertEmail,
		SendAlert: func(ip string, paths []string, dur int) error {
			if alertCalls != nil {
				atomic.AddInt64(alertCalls, 1)
			}
			return nil
		},
	})
}

func TestSensitiveThreshold(t *testing.T) {
	b := newTestBanner(2, 5, time.Minute, time.Minute, "", nil)
	b.Record("1.2.3.4", "/admin", KindSensitive)
	if b.IsBanned("1.2.3.4") {
		t.Fatal("should not be banned before reaching threshold")
	}
	b.Record("1.2.3.4", "/.env", KindSensitive)
	if !b.IsBanned("1.2.3.4") {
		t.Fatal("expected ban after reaching sensitive threshold")
	}
}

func TestNotFoundThreshold(t *testing.T) {
	b := newTestBanner(10, 3, time.Minute, time.Minute, "", nil)
	for i := 0; i < 3; i++ {
		b.Record("5.6.7.8", "/random-"+string(rune('a'+i)), KindNotFound)
	}
	if !b.IsBanned("5.6.7.8") {
		t.Fatal("expected ban after reaching 404 threshold")
	}
}

func TestNotFoundTrackDisabledByDefault(t *testing.T) {
	// 默认（TrackNotFound=false）下 404 不参与统计：正常业务路径一律放行，
	// 仅敏感扫描路径仍计数封禁——防止共享出口下正常访问被 404 探测连坐误伤。
	b := New(Options{Enabled: true, ThresholdSensitive: 1, Threshold404: 1, Window: time.Minute, Duration: time.Minute})
	for i := 0; i < 5; i++ {
		b.Record("7.7.7.7", "/random-"+string(rune('a'+i)), KindNotFound)
	}
	if b.IsBanned("7.7.7.7") {
		t.Fatal("404 events must not accumulate when TrackNotFound is false")
	}
	b.Record("7.7.7.7", "/.env", KindSensitive)
	if !b.IsBanned("7.7.7.7") {
		t.Fatal("sensitive events must still ban even when 404 tracking is off")
	}
}

func TestWindowExpiry(t *testing.T) {
	b := newTestBanner(2, 5, 50*time.Millisecond, time.Minute, "", nil)
	b.Record("1.1.1.1", "/admin", KindSensitive)
	time.Sleep(80 * time.Millisecond) // 事件滑出窗口
	b.Record("1.1.1.1", "/.env", KindSensitive)
	if b.IsBanned("1.1.1.1") {
		t.Fatal("events outside the window must not accumulate")
	}
}

func TestDurationExpiry(t *testing.T) {
	b := newTestBanner(1, 5, time.Minute, 60*time.Millisecond, "", nil)
	b.Record("2.2.2.2", "/admin", KindSensitive)
	if !b.IsBanned("2.2.2.2") {
		t.Fatal("expected ban")
	}
	time.Sleep(100 * time.Millisecond) // 封禁到期
	if b.IsBanned("2.2.2.2") {
		t.Fatal("ban should expire after duration")
	}
}

func TestAlertOncePerIP(t *testing.T) {
	var calls int64
	b := newTestBanner(1, 5, time.Minute, time.Minute, "[email protected]", &calls)
	b.Record("9.9.9.9", "/wp-admin", KindSensitive)
	b.Record("9.9.9.9", "/.git", KindSensitive) // 已在封禁中，不应重复告警
	b.Record("9.9.9.9", "/config", KindSensitive)

	// sendAlert 异步调用，轮询等待最多 500ms
	deadline := time.Now().Add(500 * time.Millisecond)
	for atomic.LoadInt64(&calls) == 0 && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if got := atomic.LoadInt64(&calls); got != 1 {
		t.Fatalf("expected exactly 1 alert, got %d", got)
	}
}

func TestDisabled(t *testing.T) {
	b := New(Options{Enabled: false})
	b.Record("3.3.3.3", "/admin", KindSensitive)
	if b.IsBanned("3.3.3.3") {
		t.Fatal("disabled auto-ban must not ban")
	}
}

func TestDefaults(t *testing.T) {
	b := New(Options{Enabled: true}) // 全部走默认值
	if b.thresholdSensitive != DefaultThresholdSensitive || b.threshold404 != DefaultThreshold404 {
		t.Fatalf("default thresholds wrong: %d/%d", b.thresholdSensitive, b.threshold404)
	}
	if b.window != DefaultWindowSeconds*time.Second || b.duration != DefaultDurationMinutes*time.Minute {
		t.Fatalf("default window/duration wrong: %v/%v", b.window, b.duration)
	}
	if b.trackNotFound {
		t.Fatal("404 tracking must be disabled by default")
	}
}

func TestIsSensitivePath(t *testing.T) {
	cases := []struct {
		path string
		want bool
	}{
		{"/admin", true},
		{"/admin/panel", true},
		{"/WP-ADMIN", true},
		{"/.env", true},
		{"/user/questions", false},
		{"/vote/abc123", false},
		{"/administratorx", false}, // 前缀必须完整段匹配，避免误伤
		// 目录枚举 / 双斜杠变体（wlwmanifest 扫描）：前缀匹配漏网，片段匹配兜底
		{"//wp/wp-includes/wlwmanifest.xml", true},
		{"//news/wp-includes/wlwmanifest.xml", true},
		{"//2018/wp-includes/wlwmanifest.xml", true},
		{"//shop/wp-includes/wlwmanifest.xml", true},
		{"//test/xmlrpc.php", true},
		{"/blog/wp-json/", true},
		{"/wp-content/uploads/evil.php", true},
		// 业务路径不得被片段匹配误伤
		{"/_/tamakyi", false},
		{"/_/admin6", false},
		{"/api/v1/question/abc123", false},
		{"/explore/tags/golang", false},
		{"/i18n/zh-CN", false},
	}
	for _, c := range cases {
		if got := IsSensitivePath(c.path); got != c.want {
			t.Errorf("IsSensitivePath(%q) = %v, want %v", c.path, got, c.want)
		}
	}
}