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, "admin@example.com", &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, "admin@example.com", &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 回调") } }