velox-bot/internal/security/gaps_test.go

package security

import (
	"encoding/json"
	"errors"
	"strings"
	"testing"
	"time"

	"velox-bot/internal/db/repos/securityrepo"

	"github.com/bwmarrin/discordgo"
)

// ---- Raid persistence ----

func TestRaidSurvivesARestart(t *testing.T) {
	e, d, st, now := testEngine(t)
	s := st.set(100)
	s.AntiRaidEnabled, s.RaidJoinCount, s.RaidAction, s.RaidQuietMinutes = true, 3, securityrepo.ActionKick, 5
	refresh(e)
	for i := 0; i < 4; i++ {
		e.OnMemberJoin(ctx, guild, d.addMember(snowflakeAt(now.AddDate(-1, 0, 0), i)))
	}
	if r, ok := st.raids[100]; !ok || !r.InvitesPaused || r.CaughtTotal != 4 || len(st.caught[100]) != 4 {
		t.Fatalf("raid not persisted: %+v caught=%d", st.raids[100], len(st.caught[100]))
	}

	// The bot restarts: a fresh engine on the same database.
	e2 := New(d, st, "999")
	e2.now = e.now
	e2.ResumeRaids(ctx)
	if !e2.RaidActive(guild) {
		t.Fatal("raid mode must resume after a restart")
	}
	// A join after the restart is still treated as a raider.
	*now = now.Add(time.Minute)
	m := d.addMember(snowflakeAt(now.AddDate(-1, 0, 0), 50))
	e2.OnMemberJoin(ctx, guild, m)
	if !contains(d.kicked, m.User.ID) {
		t.Error("raid actions continue after a restart")
	}

	// Quiet period: it ends, invites resume, the list covers both lifetimes.
	*now = now.Add(6 * time.Minute)
	e2.checkRaids(ctx)
	if e2.RaidActive(guild) {
		t.Fatal("the resumed raid should end after the quiet period")
	}
	if !strings.Contains(d.requests[len(d.requests)-1], "invites_disabled_until:<nil>") {
		t.Error("invites paused before the restart must be resumed")
	}
	var list string
	for name, content := range d.files {
		if strings.HasPrefix(name, "raid-") {
			list = content
		}
	}
	if strings.Count(list, "\tkick\t") != 5 {
		t.Errorf("summary should list all 5 raiders:\n%s", list)
	}
	if _, ok := st.raids[100]; ok {
		t.Error("the raid record should be cleared")
	}
}

// ---- Heat ----

func heatEngine(t *testing.T) (*Engine, *fakeDiscord, *fakeStore, *time.Time) {
	e, d, st, now := testEngine(t)
	s := st.set(100)
	s.HeatEnabled, s.HeatPanicMembers = true, 0
	refresh(e)
	return e, d, st, now
}

var msgN int

func say(e *Engine, userID, content string, edit ...func(m *discordgo.MessageCreate)) *discordgo.MessageCreate {
	msgN++
	m := &discordgo.MessageCreate{Message: &discordgo.Message{
		ID: "msg" + string(rune('A'+msgN%26)) + strings.Repeat("x", msgN/26), GuildID: guild, ChannelID: "300", Content: content,
		Author: &discordgo.User{ID: userID, Username: "u" + userID},
		Member: &discordgo.Member{Roles: []string{}, JoinedAt: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)},
	}}
	for _, f := range edit {
		f(m)
	}
	e.OnMessage(ctx, m)
	return m
}

func TestNormalChatNeverHeatsUp(t *testing.T) {
	e, d, _, now := heatEngine(t)
	for i := 0; i < 60; i++ {
		say(e, "50", "just talking about the game tonight, number "+string(rune('a'+i%26)))
		*now = now.Add(2 * time.Second)
	}
	if len(d.deletedMsgs) != 0 || len(d.timeouts) != 0 {
		t.Error("a message every two seconds is normal conversation")
	}
}

func TestSpamIsPunishedAndEscalates(t *testing.T) {
	e, d, st, now := heatEngine(t)
	var last *discordgo.MessageCreate
	for i := 0; i < 6 && d.timeouts["50"] == nil; i++ {
		last = say(e, "50", "BUY CHEAP STUFF NOW")
		*now = now.Add(200 * time.Millisecond)
	}
	first := d.timeouts["50"]
	if first == nil || first.Sub(*now) > 11*time.Minute || first.Sub(*now) < 9*time.Minute {
		t.Fatalf("repeated spam should earn the 10 minute timeout, got %v", first)
	}
	if !contains(d.deletedMsgs, last.ID) {
		t.Error("the message that maxed the heat should be deleted")
	}
	if !contains(st.eventKinds(), securityrepo.EventHeat) {
		t.Error("recorded")
	}

	// A second offence the same day doubles it.
	*now = now.Add(time.Hour)
	d.timeouts["50"] = nil
	for i := 0; i < 6 && d.timeouts["50"] == nil; i++ {
		say(e, "50", "BUY CHEAP STUFF NOW")
		*now = now.Add(200 * time.Millisecond)
	}
	if second := d.timeouts["50"]; second == nil || second.Sub(*now) < 19*time.Minute {
		t.Errorf("second offence should be 20 minutes, got %v", second)
	}
}

func TestInviteLinksAndMassMentions(t *testing.T) {
	e, d, _, _ := heatEngine(t)
	say(e, "51", "join my server discord.gg/abcdef")
	if d.timeouts["51"] == nil {
		t.Error("an invite link is an instant max")
	}
	say(e, "52", "hey @everyone check this")
	say(e, "52", "@everyone again")
	if d.timeouts["52"] == nil {
		t.Error("repeated @everyone should max out fast")
	}
}

func TestHeatExemptions(t *testing.T) {
	e, d, st, _ := heatEngine(t)
	st.trusted[100] = []securityrepo.Trusted{{TargetID: 60, Kind: "user"}}
	refresh(e)
	e.ChannelPerms = func(userID, _ string) (int64, error) {
		if userID == "61" {
			return discordgo.PermissionManageMessages, nil
		}
		return 0, nil
	}
	for _, who := range []string{"1", "60", "61"} {
		say(e, who, "discord.gg/x")
	}
	say(e, "62", "discord.gg/x", func(m *discordgo.MessageCreate) { m.Author.Bot = true })
	say(e, "63", "discord.gg/x", func(m *discordgo.MessageCreate) { m.WebhookID = "w" })
	if len(d.timeouts) != 0 || len(d.deletedMsgs) != 0 {
		t.Errorf("owner, trusted, moderators, bots and webhooks are exempt: %v", d.timeouts)
	}
}

func TestHeatPanicTimesOutNewMembers(t *testing.T) {
	e, d, st, now := heatEngine(t)
	st.set(100).HeatPanicMembers = 3
	refresh(e)
	for _, u := range []string{"70", "71", "72"} {
		say(e, u, "discord.gg/spam")
	}
	if !contains(st.eventKinds(), securityrepo.EventHeatPanic) {
		t.Fatal("three members maxing out within a minute starts heat panic")
	}
	newcomer := func(m *discordgo.MessageCreate) { m.Member.JoinedAt = now.Add(-time.Hour) }
	say(e, "80", "hello everyone", newcomer)
	if d.timeouts["80"] == nil {
		t.Error("someone who joined an hour ago speaking during panic is timed out")
	}
	say(e, "81", "hello everyone")
	if d.timeouts["81"] != nil {
		t.Error("established members are unaffected by panic")
	}
	*now = now.Add(11 * time.Minute)
	say(e, "82", "hello", newcomer)
	if d.timeouts["82"] != nil {
		t.Error("panic ends after 10 minutes")
	}
}

func TestMentionFloodLocksTheServer(t *testing.T) {
	e, d, st, _ := heatEngine(t)
	s := st.set(100)
	s.HeatLockdownMentions, s.HeatStrictness = 20, "low"
	refresh(e)
	d.channels["300"] = &discordgo.Channel{ID: "300", GuildID: guild, Type: discordgo.ChannelTypeGuildText}
	five := func(m *discordgo.MessageCreate) {
		m.Mentions = []*discordgo.User{{ID: "a"}, {ID: "b"}, {ID: "c"}, {ID: "d"}, {ID: "e"}}
	}
	for i, u := range []string{"90", "91", "92", "93"} {
		_ = i
		say(e, u, "hi", five)
	}
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		st.mu.Lock()
		n := len(st.lockdowns)
		st.mu.Unlock()
		if n > 0 {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Error("20 mentions within 30 seconds should lock the server")
}

func TestHeatAlertsAreCapped(t *testing.T) {
	e, d, _, _ := heatEngine(t)
	for i := 0; i < 20; i++ {
		say(e, string(rune('A'+i)), "discord.gg/x")
	}
	if n := d.sentCount(); n > heatAlertsPerMin {
		t.Errorf("a spam wave produced %d alerts, cap is %d", n, heatAlertsPerMin)
	}
}

// ---- Snapshots and restore ----

func TestSnapshotAndRestore(t *testing.T) {
	e, d, st, _ := testEngine(t)
	d.channels["400"] = &discordgo.Channel{ID: "400", GuildID: guild, Type: discordgo.ChannelTypeGuildCategory, Name: "Staff"}
	d.channels["401"] = &discordgo.Channel{ID: "401", GuildID: guild, Type: discordgo.ChannelTypeGuildText, Name: "mod-chat", ParentID: "400",
		PermissionOverwrites: []*discordgo.PermissionOverwrite{
			{ID: guild, Type: discordgo.PermissionOverwriteTypeRole, Deny: discordgo.PermissionViewChannel},
			{ID: "201", Type: discordgo.PermissionOverwriteTypeRole, Allow: discordgo.PermissionViewChannel},
		}}
	d.channels["402"] = &discordgo.Channel{ID: "402", GuildID: guild, Type: discordgo.ChannelTypeGuildPublicThread, Name: "a thread"}

	if _, err := e.Restore(ctx, guild, "1"); !errors.As(err, new(UserError)) {
		t.Fatalf("no snapshot yet should be a clear message: %v", err)
	}
	roles, channels, err := e.TakeSnapshot(ctx, guild, "1")
	if err != nil || roles != 5 || channels != 2 {
		t.Fatalf("snapshot = %d roles, %d channels, %v (managed roles and threads are skipped)", roles, channels, err)
	}
	var data snapshotData
	_ = json.Unmarshal(st.snapshots[0].Data, &data)
	for _, r := range data.Roles {
		if r.ID == "204" || r.ID == "299" {
			t.Error("integration roles must not be in the snapshot")
		}
	}

	// The nuke: the Mod role and the whole Staff category are deleted, and
	// the Member role is given Administrator.
	var keep []*discordgo.Role
	for _, r := range d.roles {
		if r.ID != "201" {
			keep = append(keep, r)
		}
		if r.ID == "202" {
			r.Permissions |= discordgo.PermissionAdministrator
		}
	}
	d.roles = keep
	delete(d.channels, "400")
	delete(d.channels, "401")

	res, err := e.Restore(ctx, guild, "1")
	if err != nil {
		t.Fatal(err)
	}
	if res.RolesCreated != 1 || res.RolesReverted != 1 || res.ChannelsCreated != 2 || res.Failed != 0 {
		t.Fatalf("restore result = %+v", res)
	}
	var newMod, newCat, newText *string
	for _, r := range d.roles {
		if r.Name == "Mod" {
			newMod = &r.ID
		}
		if r.ID == "202" && r.Permissions&discordgo.PermissionAdministrator != 0 {
			t.Error("Administrator added to Member should be reverted")
		}
	}
	for _, c := range d.channels {
		switch c.Name {
		case "Staff":
			newCat = &c.ID
		case "mod-chat":
			newText = &c.ID
		}
	}
	if newMod == nil || newCat == nil || newText == nil {
		t.Fatal("role or channels not recreated")
	}
	text := d.channels[*newText]
	if text.ParentID != *newCat {
		t.Error("the channel should go back into the recreated category")
	}
	foundMod := false
	for _, o := range text.PermissionOverwrites {
		if o.ID == *newMod && o.Allow == discordgo.PermissionViewChannel {
			foundMod = true
		}
		if o.ID == "201" {
			t.Error("overwrites must point at the recreated role, not the deleted one")
		}
	}
	if !foundMod {
		t.Errorf("the Mod overwrite should be restored: %+v", text.PermissionOverwrites)
	}

	// Restoring again finds nothing missing.
	res, _ = e.Restore(ctx, guild, "1")
	if res.RolesCreated+res.ChannelsCreated+res.RolesReverted != 0 {
		t.Errorf("a second restore should change nothing: %+v", res)
	}
}

func TestOnlyOneRestoreAtATime(t *testing.T) {
	e, _, _, _ := testEngine(t)
	e.restoring[guild] = true
	if _, err := e.Restore(ctx, guild, "1"); err == nil || !strings.Contains(err.Error(), "already running") {
		t.Errorf("err = %v", err)
	}
}