velox-bot/internal/security/heat.go

package security

import (
	"context"
	"fmt"
	"log"
	"math"
	"regexp"
	"strings"
	"sync"
	"time"

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

	"github.com/bwmarrin/discordgo"
)

// Heat is a per-member spam score, like Wick's. Every message adds heat,
// more for what spam looks like; heat cools down over time, so normal
// conversation never builds up. At maxHeat the message is deleted and the
// configured action applied.
const (
	maxHeat          = 100.0
	heatDecayPerSec  = 8.0 // a message every couple of seconds never accumulates
	heatBase         = 8.0
	heatRepeat       = 25.0
	heatMention      = 10.0
	heatEveryone     = 60.0
	heatEmoji        = 2.0
	heatLine         = 2.0 // per line beyond the fourth
	heatPer100Chars  = 1.0
	heatAttachment   = 10.0
	heatLink         = 8.0
	heatBlockedWord  = 60.0
	maxHeatUsers     = 50000 // per guild, swept when exceeded
	panicWindow      = time.Minute
	panicDuration    = 10 * time.Minute
	panicTimeout     = time.Hour
	panicNewMember   = 24 * time.Hour
	mentionWindow    = 30 * time.Second
	floodCooldown    = 10 * time.Minute
	heatAlertsPerMin = 5
	maxTimeout       = 28*24*time.Hour - time.Hour
)

var (
	linkRe        = regexp.MustCompile(`https?://\S+`)
	customEmojiRe = regexp.MustCompile(`<a?:\w+:\d+>`)
)

var strictness = map[string]float64{"low": 0.6, "medium": 1.0, "high": 1.5}

type memberHeat struct {
	heat        float64
	last        time.Time
	recent      []string // last few messages, normalized, for repeats
	offences    int      // within the last day
	lastOffence time.Time
}

type guildHeat struct {
	mu          sync.Mutex
	members     map[string]*memberHeat
	maxedAt     map[string]time.Time // who hit max heat recently, for panic
	panicUntil  time.Time
	mentions    []time.Time // mentions by non-exempt members, for flood lockdown
	lastFlood   time.Time
	alertTimes  []time.Time
	alertsMuted int
}

func (e *Engine) heatFor(guildID string) *guildHeat {
	e.mu.Lock()
	defer e.mu.Unlock()
	g, ok := e.heat[guildID]
	if !ok {
		g = &guildHeat{members: map[string]*memberHeat{}, maxedAt: map[string]time.Time{}}
		e.heat[guildID] = g
	}
	return g
}

func normalize(s string) string {
	return strings.Join(strings.Fields(strings.ToLower(s)), " ")
}

func countEmoji(s string) int {
	n := len(customEmojiRe.FindAllString(s, -1))
	for _, r := range s {
		if r >= 0x1F000 || (r >= 0x2600 && r <= 0x27BF) {
			n++
		}
	}
	return n
}

// MessageHeat scores one message. recent is the author's last few messages
// (normalized) for spotting repeats. The second result is true when the
// message is an instant maximum (an invite link).
func MessageHeat(s *securityrepo.Settings, m *discordgo.Message, recent []string) (float64, bool, string) {
	content := m.Content
	lower := strings.ToLower(content)
	if inviteRe.MatchString(lower) {
		return maxHeat, true, "invite link"
	}
	reason := "spam"
	h := heatBase
	norm := normalize(content)
	if norm != "" {
		for _, r := range recent {
			if r == norm {
				h += heatRepeat
				reason = "repeated messages"
				break
			}
		}
	}
	mentions := len(m.Mentions) + len(m.MentionRoles)
	h += float64(mentions) * heatMention
	if m.MentionEveryone || strings.Contains(lower, "@everyone") || strings.Contains(lower, "@here") {
		h += heatEveryone
		reason = "mass mentions"
	} else if mentions >= 3 {
		reason = "mention spam"
	}
	h += float64(countEmoji(content)) * heatEmoji
	if lines := strings.Count(content, "\n") + 1; lines > 4 {
		h += float64(lines-4) * heatLine
	}
	h += float64(len([]rune(content))/100) * heatPer100Chars
	h += float64(len(m.Attachments)+len(m.StickerItems)) * heatAttachment
	h += float64(len(linkRe.FindAllString(content, -1))) * heatLink
	for _, w := range s.HeatBlockedWords {
		if w = strings.ToLower(strings.TrimSpace(w)); w != "" && strings.Contains(lower, w) {
			h += heatBlockedWord
			reason = "blocked word"
			break
		}
	}
	mult := strictness[s.HeatStrictness]
	if mult == 0 {
		mult = 1
	}
	return h * mult, false, reason
}

// heatExempt: bots, webhooks, the owner, trusted members and anyone who can
// manage messages in the channel.
func (e *Engine) heatExempt(ctx context.Context, m *discordgo.MessageCreate) bool {
	if m.Author == nil || m.Author.Bot || m.WebhookID != "" {
		return true
	}
	var roles []string
	if m.Member != nil {
		roles = m.Member.Roles
	}
	if e.IsTrusted(ctx, m.GuildID, m.Author.ID, nonNil(roles)) {
		return true
	}
	if e.ChannelPerms != nil {
		if p, err := e.ChannelPerms(m.Author.ID, m.ChannelID); err == nil &&
			p&(discordgo.PermissionAdministrator|discordgo.PermissionManageMessages) != 0 {
			return true
		}
	}
	return false
}

// OnMessage runs the heat system on a guild message.
func (e *Engine) OnMessage(ctx context.Context, m *discordgo.MessageCreate) {
	if m == nil || m.Message == nil || m.GuildID == "" {
		return
	}
	s, _, err := e.load(ctx, m.GuildID)
	if err != nil || !s.HeatEnabled {
		return
	}
	if e.heatExempt(ctx, m) {
		return
	}
	now := e.now()
	g := e.heatFor(m.GuildID)

	// Mention flood: lock the server once, then cool down.
	if s.HeatLockdownMentions > 0 {
		n := len(m.Mentions) + len(m.MentionRoles)
		if m.MentionEveryone {
			n += 10
		}
		if n > 0 && e.mentionFlood(g, now, n, s.HeatLockdownMentions) {
			go e.floodLockdown(context.Background(), m.GuildID, s)
		}
	}

	g.mu.Lock()
	if now.Before(g.panicUntil) && m.Member != nil && !m.Member.JoinedAt.IsZero() && now.Sub(m.Member.JoinedAt) < panicNewMember {
		g.mu.Unlock()
		e.heatPunish(ctx, m, s, "heat panic: new member speaking during a spam wave", panicTimeout, securityrepo.ActionTimeout)
		return
	}
	mh, ok := g.members[m.Author.ID]
	if !ok {
		if len(g.members) >= maxHeatUsers {
			g.sweep(now)
		}
		mh = &memberHeat{}
		g.members[m.Author.ID] = mh
	}
	if !mh.last.IsZero() {
		mh.heat = math.Max(0, mh.heat-now.Sub(mh.last).Seconds()*heatDecayPerSec)
	}
	mh.last = now
	add, instant, reason := MessageHeat(s, m.Message, mh.recent)
	mh.heat += add
	if norm := normalize(m.Content); norm != "" {
		mh.recent = append(mh.recent, norm)
		if len(mh.recent) > 3 {
			mh.recent = mh.recent[len(mh.recent)-3:]
		}
	}
	if mh.heat < maxHeat && !instant {
		g.mu.Unlock()
		return
	}
	// Over the top: reset, count the offence, and work out the timeout.
	mh.heat = 0
	if now.Sub(mh.lastOffence) > 24*time.Hour {
		mh.offences = 0
	}
	mh.offences++
	mh.lastOffence = now
	timeout := time.Duration(s.HeatTimeoutMinutes) * time.Minute
	if s.HeatEscalate && mh.offences > 1 {
		timeout *= time.Duration(1 << min(mh.offences-1, 12))
	}
	if timeout > maxTimeout {
		timeout = maxTimeout
	}
	// Heat panic: enough different members maxing out within a minute.
	g.maxedAt[m.Author.ID] = now
	for u, t := range g.maxedAt {
		if now.Sub(t) > panicWindow {
			delete(g.maxedAt, u)
		}
	}
	startPanic := s.HeatPanicMembers > 0 && len(g.maxedAt) >= s.HeatPanicMembers && !now.Before(g.panicUntil)
	if startPanic {
		g.panicUntil = now.Add(panicDuration)
	}
	g.mu.Unlock()

	e.heatPunish(ctx, m, s, reason, timeout, s.HeatAction)
	if startPanic {
		e.event(ctx, m.GuildID, securityrepo.EventHeatPanic, e.botID, "", fmt.Sprintf("%d members hit max heat within a minute", s.HeatPanicMembers))
		e.send(ctx, m.GuildID, s, alert{ping: true, critical: true, embed: &discordgo.MessageEmbed{
			Title:       "🔥 Spam wave: heat panic on",
			Description: fmt.Sprintf("%d members hit max heat within a minute. For the next 10 minutes, anyone who joined in the last day is timed out as soon as they speak.", s.HeatPanicMembers),
			Color:       colorAlert,
		}})
	}
}

func (g *guildHeat) sweep(now time.Time) {
	for u, mh := range g.members {
		cooled := mh.heat-now.Sub(mh.last).Seconds()*heatDecayPerSec <= 0
		if cooled && now.Sub(mh.lastOffence) > 24*time.Hour {
			delete(g.members, u)
		}
	}
}

// mentionFlood records n mentions and reports whether the guild just went
// over the limit (once per cooldown).
func (e *Engine) mentionFlood(g *guildHeat, now time.Time, n, limit int) bool {
	g.mu.Lock()
	defer g.mu.Unlock()
	for i := 0; i < n && len(g.mentions) < 5000; i++ {
		g.mentions = append(g.mentions, now)
	}
	cut := 0
	for cut < len(g.mentions) && now.Sub(g.mentions[cut]) > mentionWindow {
		cut++
	}
	g.mentions = g.mentions[cut:]
	if len(g.mentions) >= limit && now.Sub(g.lastFlood) > floodCooldown {
		g.lastFlood = now
		g.mentions = nil
		return true
	}
	return false
}

func (e *Engine) floodLockdown(ctx context.Context, guildID string, s *securityrepo.Settings) {
	locked, failed, err := e.LockServer(ctx, guildID, e.botID)
	desc := fmt.Sprintf("%d or more mentions within 30 seconds: every channel is locked (%d locked", s.HeatLockdownMentions, locked)
	if failed > 0 {
		desc += fmt.Sprintf(", %d failed", failed)
	}
	desc += "). Unlock with `/unlock server` once it's under control."
	if err != nil {
		desc = "A mention flood was detected but locking the server failed: " + clip(err.Error(), 200)
	}
	e.send(ctx, guildID, s, alert{ping: true, critical: true, embed: &discordgo.MessageEmbed{Title: "🔒 Mention flood: server locked", Description: desc, Color: colorAlert}})
}

// heatPunish deletes the message and applies the action. Alerts are capped
// per minute so a spam wave can't flood the alert channel.
func (e *Engine) heatPunish(ctx context.Context, m *discordgo.MessageCreate, s *securityrepo.Settings, reason string, timeout time.Duration, action string) {
	if err := e.d.ChannelMessageDelete(m.ChannelID, m.ID, discordgo.WithAuditLogReason("Velox heat: "+reason)); err != nil {
		log.Printf("security: deleting a hot message in guild %s: %v", m.GuildID, err)
	}
	var err error
	if action == securityrepo.ActionTimeout {
		if m.Author.ID == e.botID || e.IsOwner(m.GuildID, m.Author.ID) {
			return
		}
		until := e.now().Add(timeout)
		err = e.d.GuildMemberTimeout(m.GuildID, m.Author.ID, &until, discordgo.WithAuditLogReason(clip("Velox heat: "+reason, 400)))
	} else {
		err = e.punish(m.GuildID, m.Author, action, "heat: "+reason)
	}
	if err != nil {
		log.Printf("security: heat action on %s in guild %s: %v", m.Author.ID, m.GuildID, err)
		return
	}
	what := action
	if action == securityrepo.ActionTimeout {
		what = "timed out for " + humanize(timeout)
	}
	e.event(ctx, m.GuildID, securityrepo.EventHeat, e.botID, m.Author.ID, what+": "+reason)

	g := e.heatFor(m.GuildID)
	now := e.now()
	g.mu.Lock()
	cut := 0
	for cut < len(g.alertTimes) && now.Sub(g.alertTimes[cut]) > time.Minute {
		cut++
	}
	g.alertTimes = g.alertTimes[cut:]
	allowed := len(g.alertTimes) < heatAlertsPerMin
	muted := 0
	if allowed {
		g.alertTimes = append(g.alertTimes, now)
		muted, g.alertsMuted = g.alertsMuted, 0
	} else {
		g.alertsMuted++
	}
	g.mu.Unlock()
	if !allowed {
		return
	}
	desc := fmt.Sprintf("<@%s> (%s) was %s: %s. Their message was deleted.", m.Author.ID, safeName(m.Author), what, reason)
	if muted > 0 {
		desc += fmt.Sprintf("\n(%d more since the last alert weren't posted, to keep this channel readable.)", muted)
	}
	e.send(ctx, m.GuildID, s, alert{embed: &discordgo.MessageEmbed{Title: "Heat", Description: desc, Color: colorWarn}})
}

func humanize(d time.Duration) string {
	switch {
	case d >= 24*time.Hour:
		return fmt.Sprintf("%.0f days", d.Hours()/24)
	case d >= time.Hour:
		return fmt.Sprintf("%.0f hours", d.Hours())
	default:
		return fmt.Sprintf("%.0f minutes", math.Max(1, d.Minutes()))
	}
}