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)
	}
}

// TestAlertCooldownExpires 回归「告警冷却期」而非「一生只发一次」：
// 冷却期（alertCooldown = 24h）过后，同一来源再次触发封禁应能再告警。
//
// 早先是 map[string]bool 永久去重，长期盯防同一来源会漏掉后续攻击活动
// （换路径再来、隔几天再来都收不到提醒）。这里直接回拨 alerted 的时间戳，
// 模拟「冷却期已过」，断言能再次发出告警。
func TestAlertCooldownExpires(t *testing.T) {
	var calls int64
	b := newTestBanner(1, 5, time.Minute, time.Minute, "[email protected]", &calls)

	// 第一次封禁 + 告警
	b.Record("8.8.8.8", "/admin", KindSensitive)
	waitFor(&calls, 1, t)

	// 冷却期内再次命中：不应重复告警（此时已在封禁中，先手动解除封禁以走到告警判定）
	b.mu.Lock()
	delete(b.banned, "8.8.8.8")
	b.mu.Unlock()
	b.Record("8.8.8.8", "/config", KindSensitive)
	time.Sleep(50 * time.Millisecond)
	if got := atomic.LoadInt64(&calls); got != 1 {
		t.Fatalf("冷却期内不应重复告警，got %d", got)
	}

	// 回拨 alerted 时间戳到冷却期之前，模拟「24 小时后攻击者再来」
	b.mu.Lock()
	b.alerted["8.8.8.8"] = time.Now().Add(-alertCooldown - time.Minute)
	delete(b.banned, "8.8.8.8")
	b.mu.Unlock()
	b.Record("8.8.8.8", "/.env", KindSensitive)

	waitFor(&calls, 2, t)
	if got := atomic.LoadInt64(&calls); got != 2 {
		t.Fatalf("冷却期过后应再次告警，got %d", got)
	}
}

// waitFor 轮询等待 atomic 计数器达到 want（sendAlert 是异步 goroutine）。
func waitFor(c *int64, want int64, t *testing.T) {
	t.Helper()
	deadline := time.Now().Add(500 * time.Millisecond)
	for atomic.LoadInt64(c) < want && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
}

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)
		}
	}
}

// TestOnBanCallback 验证建封时会回调上层（键 / 时长 / 原因），
// 且「已在封禁期内重复命中」不会重复回调。
//
// 存在的意义：封禁表是纯内存的，外部完全看不见。线上出现过
// 「邮件说封了、实际还能访问」，而当时没有任何记录能说明封的是哪个键、
// 到什么时候解除 —— 这个回调就是那条记录的唯一来源。
func TestOnBanCallback(t *testing.T) {
	type banEvent struct {
		key      string
		duration time.Duration
		reason   string
	}
	events := make(chan banEvent, 4)
	b := New(Options{
		Enabled:            true,
		ThresholdSensitive: 1,
		TrackNotFound:      false,
		Window:             time.Minute,
		Duration:           10 * time.Minute,
		HoneypotEnabled:    true,
		HoneypotThreshold:  1,
		HoneypotDuration:   24 * time.Hour,
		OnBan: func(key string, duration time.Duration, reason string) {
			events <- banEvent{key: key, duration: duration, reason: reason}
		},
	})

	const key = "2408:8207:1234:5678::/64"
	if !b.RecordHoneypot(key, "/wp-login.php") {
		t.Fatal("首次蜜罐命中应触发封禁")
	}
	select {
	case e := <-events:
		if e.key != key {
			t.Errorf("回调键 = %q, want %q", e.key, key)
		}
		if e.duration != 24*time.Hour {
			t.Errorf("回调时长 = %v, want 24h（蜜罐应走 HoneypotDuration）", e.duration)
		}
		if e.reason != ReasonHoneypot {
			t.Errorf("回调原因 = %q, want %q", e.reason, ReasonHoneypot)
		}
	case <-time.After(time.Second):
		t.Fatal("蜜罐封禁没有触发 OnBan 回调")
	}

	// 已在封禁期内再次命中：不新建封禁，也就不应重复回调。
	b.RecordHoneypot(key, "/wp-login.php")
	select {
	case e := <-events:
		t.Fatalf("重复命中不应重复回调，收到 %+v", e)
	case <-time.After(50 * time.Millisecond):
	}

	// 常规通道（敏感路径）也要回调，且时长用常规 Duration、原因可区分。
	b.Record("1.2.3.4", "/admin", KindSensitive)
	select {
	case e := <-events:
		if e.key != "1.2.3.4" || e.reason != ReasonProbe || e.duration != 10*time.Minute {
			t.Errorf("常规封禁回调 = %+v, want key=1.2.3.4 reason=probe duration=10m", e)
		}
	case <-time.After(time.Second):
		t.Fatal("常规封禁没有触发 OnBan 回调")
	}
}