velox-bot/internal/security/nuke.go

package security

import (
	"context"
	"fmt"
	"log"
	"strconv"
	"sync"
	"time"

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

	"github.com/bwmarrin/discordgo"
)

// nukeCounter counts destructive actions per server, person and kind over a
// sliding window, and remembers whom it just punished so one nuke attempt
// is answered once.
type nukeCounter struct {
	mu       sync.Mutex
	hits     map[string][]time.Time
	punished map[string]time.Time
}

const (
	maxHitsPerKey  = 200
	punishCooldown = time.Minute
)

func newNukeCounter() *nukeCounter {
	return &nukeCounter{hits: map[string][]time.Time{}, punished: map[string]time.Time{}}
}

// add records a hit and returns how many fall inside the window.
func (c *nukeCounter) add(key string, now time.Time, window time.Duration) int {
	c.mu.Lock()
	defer c.mu.Unlock()
	h := append(c.hits[key], now)
	cut := 0
	for cut < len(h) && now.Sub(h[cut]) > window {
		cut++
	}
	h = h[cut:]
	if len(h) > maxHitsPerKey {
		h = h[len(h)-maxHitsPerKey:]
	}
	c.hits[key] = h
	// Opportunistic cleanup so the map can't grow forever.
	if len(c.hits) > 5000 {
		for k, v := range c.hits {
			if len(v) == 0 || now.Sub(v[len(v)-1]) > window {
				delete(c.hits, k)
			}
		}
	}
	return len(h)
}

// claim reports whether this person may be punished now (not punished in the
// last minute), and marks them.
func (c *nukeCounter) claim(key string, now time.Time) bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	if t, ok := c.punished[key]; ok && now.Sub(t) < punishCooldown {
		return false
	}
	c.punished[key] = now
	if len(c.punished) > 5000 {
		for k, t := range c.punished {
			if now.Sub(t) > punishCooldown {
				delete(c.punished, k)
			}
		}
	}
	return true
}

func changeKey(c *discordgo.AuditLogChange) string {
	if c == nil || c.Key == nil {
		return ""
	}
	return string(*c.Key)
}

func permValue(v any) int64 {
	switch x := v.(type) {
	case string:
		n, _ := strconv.ParseInt(x, 10, 64)
		return n
	case float64:
		return int64(x)
	}
	return 0
}

// Classify maps an audit log entry to a limit key, or flags it as always
// forbidden for untrusted members (instant). rolePerms looks up a role's
// permissions for role grants.
func Classify(entry *discordgo.AuditLogEntry, rolePerms func(roleID string) int64) (key string, instant bool, what string) {
	if entry == nil || entry.ActionType == nil {
		return "", false, ""
	}
	switch *entry.ActionType {
	case discordgo.AuditLogActionChannelCreate:
		return securityrepo.LimitChannelCreate, false, "created channels"
	case discordgo.AuditLogActionChannelDelete:
		return securityrepo.LimitChannelDelete, false, "deleted channels"
	case discordgo.AuditLogActionRoleDelete:
		return securityrepo.LimitRoleDelete, false, "deleted roles"
	case discordgo.AuditLogActionMemberBanAdd:
		return securityrepo.LimitBan, false, "banned members"
	case discordgo.AuditLogActionMemberKick:
		return securityrepo.LimitKick, false, "kicked members"
	case discordgo.AuditLogActionWebhookCreate:
		return securityrepo.LimitWebhookCreate, false, "created webhooks"
	case discordgo.AuditLogActionMemberPrune:
		return "", true, "pruned members"
	case discordgo.AuditLogActionRoleCreate:
		for _, c := range entry.Changes {
			if changeKey(c) == "permissions" && permValue(c.NewValue)&DangerousPerms != 0 {
				return "", true, "created a role with dangerous permissions"
			}
		}
		return securityrepo.LimitRoleCreate, false, "created roles"
	case discordgo.AuditLogActionRoleUpdate:
		for _, c := range entry.Changes {
			if changeKey(c) == "permissions" {
				added := permValue(c.NewValue) &^ permValue(c.OldValue)
				if added&DangerousPerms != 0 {
					return "", true, "gave a role dangerous permissions"
				}
			}
		}
	case discordgo.AuditLogActionMemberRoleUpdate:
		for _, c := range entry.Changes {
			if changeKey(c) != "$add" {
				continue
			}
			list, _ := c.NewValue.([]any)
			for _, item := range list {
				if r, ok := item.(map[string]any); ok {
					if rid, _ := r["id"].(string); rid != "" && rolePerms != nil && rolePerms(rid)&DangerousPerms != 0 {
						return "", true, "gave someone a role with dangerous permissions"
					}
				}
			}
		}
	case discordgo.AuditLogActionGuildUpdate:
		for _, c := range entry.Changes {
			if changeKey(c) == "vanity_url_code" {
				return "", true, "changed the vanity URL"
			}
		}
	}
	return "", false, ""
}

// OnAuditLogEntry watches what members do: bots added by untrusted members
// (Join Gate) and destructive actions over the limits (anti-nuke).
func (e *Engine) OnAuditLogEntry(ctx context.Context, guildID string, entry *discordgo.AuditLogEntry) {
	if entry == nil || entry.ActionType == nil || entry.UserID == "" {
		return
	}
	s, _, err := e.load(ctx, guildID)
	if err != nil {
		log.Printf("security: settings for guild %s: %v", guildID, err)
		return
	}

	if *entry.ActionType == discordgo.AuditLogActionBotAdd {
		e.checkBotAdd(ctx, guildID, s, entry)
		return
	}
	if !s.AntiNukeEnabled || !worthChecking(entry) {
		return
	}

	// Role grants need care: someone could give themselves (or an alt) a
	// trusted role to slip past everything below. Roles added by this very
	// entry never count towards the actor's own trust, and an untrusted
	// member granting a trusted role is stopped and the role taken back.
	added := addedRoles(entry)
	var actorRoles []string
	if m, err := e.d.GuildMember(guildID, entry.UserID); err == nil {
		actorRoles = m.Roles
		if entry.TargetID == entry.UserID {
			actorRoles = without(actorRoles, added)
		}
	}
	if e.IsTrusted(ctx, guildID, entry.UserID, nonNil(actorRoles)) {
		return
	}

	key, instant, what := Classify(entry, func(roleID string) int64 { return e.rolePerms(guildID, roleID) })
	if granted := e.trustedAmong(ctx, guildID, added); len(granted) > 0 {
		instant, what = true, "gave someone a trusted role"
		e.revokeRoles(guildID, entry.TargetID, granted)
	}
	if *entry.ActionType == discordgo.AuditLogActionRoleUpdate && e.isBotRole(guildID, entry.TargetID) {
		instant, what = true, "changed Velox's own role"
	}
	if key == "" && !instant {
		return
	}
	who := guildID + "|" + entry.UserID
	if !instant {
		limit, ok := s.AntiNukeLimits[key]
		if !ok {
			limit = securityrepo.DefaultLimits()[key]
		}
		n := e.nukes.add(who+"|"+key, e.now(), time.Duration(s.AntiNukeWindowSeconds)*time.Second)
		if n <= limit {
			return
		}
		what = fmt.Sprintf("%s %d times in %d seconds (limit %d)", what, n, s.AntiNukeWindowSeconds, limit)
	}
	if !e.nukes.claim(who, e.now()) {
		return
	}
	e.stopNuke(ctx, guildID, s, entry.UserID, what)
}

func (e *Engine) rolePerms(guildID, roleID string) int64 {
	for _, r := range e.guildRoles(guildID) {
		if r.ID == roleID {
			return r.Permissions
		}
	}
	return 0
}

// worthChecking skips audit log entries that can never trigger anti-nuke
// before anything is looked up on Discord. A busy server writes many entries
// (message deletes, nickname changes...), and lookups for each would eat
// into the rate limit the bot needs most during an attack.
func worthChecking(entry *discordgo.AuditLogEntry) bool {
	switch *entry.ActionType {
	case discordgo.AuditLogActionChannelCreate, discordgo.AuditLogActionChannelDelete,
		discordgo.AuditLogActionRoleCreate, discordgo.AuditLogActionRoleDelete, discordgo.AuditLogActionRoleUpdate,
		discordgo.AuditLogActionMemberBanAdd, discordgo.AuditLogActionMemberKick, discordgo.AuditLogActionMemberPrune,
		discordgo.AuditLogActionWebhookCreate:
		return true
	case discordgo.AuditLogActionMemberRoleUpdate:
		return len(addedRoles(entry)) > 0 // removing roles is harmless
	case discordgo.AuditLogActionGuildUpdate:
		for _, c := range entry.Changes {
			if changeKey(c) == "vanity_url_code" {
				return true
			}
		}
	}
	return false
}

const roleCacheTTL = 10 * time.Second

type cachedRoles struct {
	roles []*discordgo.Role
	at    time.Time
}

// guildRoles returns the guild's roles, cached briefly. Staleness can't hide
// an escalation: creating or editing a role with dangerous permissions is
// caught from its own audit log entry.
func (e *Engine) guildRoles(guildID string) []*discordgo.Role {
	e.mu.Lock()
	c, ok := e.roleCache[guildID]
	e.mu.Unlock()
	if ok && e.now().Sub(c.at) < roleCacheTTL {
		return c.roles
	}
	roles, err := e.d.GuildRoles(guildID)
	if err != nil {
		return nil
	}
	e.mu.Lock()
	if e.roleCache == nil {
		e.roleCache = map[string]cachedRoles{}
	}
	e.roleCache[guildID] = cachedRoles{roles: roles, at: e.now()}
	e.mu.Unlock()
	return roles
}

func (e *Engine) stopNuke(ctx context.Context, guildID string, s *securityrepo.Settings, userID, what string) {
	reason := "anti-nuke: " + what
	e.event(ctx, guildID, securityrepo.EventNuke, userID, userID, what)

	desc := fmt.Sprintf("<@%s> %s.", userID, what)
	if err := e.Quarantine(ctx, guildID, userID, reason); err != nil {
		log.Printf("security: quarantining %s in guild %s: %v", userID, guildID, err)
		desc += "\n⚠️ I couldn't fully quarantine them: " + clip(err.Error(), 300) +
			"\nMove the Velox role above their highest role and give it Manage Roles, Moderate Members and Kick Members."
	} else {
		desc += "\nThey were quarantined: roles removed (saved) and timed out. Release them with `/security release` if this was a mistake."
	}
	if s.AntiNukeLockServer {
		locked, failed, err := e.LockServer(ctx, guildID, e.botID)
		switch {
		case err != nil:
			desc += "\n⚠️ Locking the server failed: " + clip(err.Error(), 200)
		default:
			desc += fmt.Sprintf("\nThe server is locked (%d channels", locked)
			if failed > 0 {
				desc += fmt.Sprintf(", %d couldn't be locked", failed)
			}
			desc += "). Unlock it with `/unlock server`."
		}
	}
	if s.AntiNukeAutoRestore {
		desc += "\nRestoring deleted roles and channels from the latest snapshot; a report follows."
		go func() {
			// Give the attack a moment to stop, then rebuild.
			time.Sleep(5 * time.Second)
			if _, err := e.Restore(context.Background(), guildID, e.botID); err != nil {
				log.Printf("security: auto restore in guild %s: %v", guildID, err)
			}
		}()
	}
	e.send(ctx, guildID, s, alert{
		ping: true, pingOwner: true, critical: true,
		embed: &discordgo.MessageEmbed{
			Title: "🛡️ Nuke attempt stopped", Description: desc, Color: colorAlert,
			Timestamp: e.now().UTC().Format(time.RFC3339),
		},
	})
}

// checkBotAdd removes a bot that someone other than the owner or a trusted
// member added, when the Join Gate is on.
func (e *Engine) checkBotAdd(ctx context.Context, guildID string, s *securityrepo.Settings, entry *discordgo.AuditLogEntry) {
	if !s.JoinGateEnabled || s.UnauthorizedBot == securityrepo.ActionOff || entry.TargetID == "" {
		return
	}
	if e.IsTrusted(ctx, guildID, entry.UserID, nil) {
		return
	}
	bot := &discordgo.User{ID: entry.TargetID}
	if err := e.punish(guildID, bot, s.UnauthorizedBot, fmt.Sprintf("bot added by %s, who isn't trusted", entry.UserID)); err != nil {
		log.Printf("security: removing unauthorized bot %s in guild %s: %v", entry.TargetID, guildID, err)
		return
	}
	e.event(ctx, guildID, securityrepo.EventUnauthorizedBot, entry.UserID, entry.TargetID, s.UnauthorizedBot)
	e.send(ctx, guildID, s, alert{
		ping: true, critical: true,
		embed: &discordgo.MessageEmbed{
			Title:       "Bot removed",
			Description: fmt.Sprintf("<@%s> added the bot <@%s> but isn't on the trusted list, so the bot was %s.", entry.UserID, entry.TargetID, map[string]string{securityrepo.ActionKick: "kicked", securityrepo.ActionBan: "banned"}[s.UnauthorizedBot]),
			Color:       colorWarn,
		},
	})
}

// addedRoles lists the role IDs a member role update added.
func addedRoles(entry *discordgo.AuditLogEntry) []string {
	if entry.ActionType == nil || *entry.ActionType != discordgo.AuditLogActionMemberRoleUpdate {
		return nil
	}
	var out []string
	for _, c := range entry.Changes {
		if changeKey(c) != "$add" {
			continue
		}
		list, _ := c.NewValue.([]any)
		for _, item := range list {
			if r, ok := item.(map[string]any); ok {
				if rid, _ := r["id"].(string); rid != "" {
					out = append(out, rid)
				}
			}
		}
	}
	return out
}

func without(list, remove []string) []string {
	drop := map[string]bool{}
	for _, r := range remove {
		drop[r] = true
	}
	out := []string{}
	for _, r := range list {
		if !drop[r] {
			out = append(out, r)
		}
	}
	return out
}

// nonNil keeps an empty role list from meaning "look the member up".
func nonNil(roles []string) []string {
	if roles == nil {
		return []string{}
	}
	return roles
}

// trustedAmong returns which of the roles are on the trusted list.
func (e *Engine) trustedAmong(ctx context.Context, guildID string, roles []string) []string {
	if len(roles) == 0 {
		return nil
	}
	_, trusted, err := e.load(ctx, guildID)
	if err != nil {
		return nil
	}
	var out []string
	for _, r := range roles {
		for _, t := range trusted {
			if t.Kind == "role" && id(t.TargetID) == r {
				out = append(out, r)
			}
		}
	}
	return out
}

// revokeRoles takes roles back from a member (best effort; never the owner).
func (e *Engine) revokeRoles(guildID, userID string, roles []string) {
	if userID == "" || e.IsOwner(guildID, userID) {
		return
	}
	m, err := e.d.GuildMember(guildID, userID)
	if err != nil {
		return
	}
	next := without(m.Roles, roles)
	if _, err := e.d.GuildMemberEdit(guildID, userID, &discordgo.GuildMemberParams{Roles: &next},
		discordgo.WithAuditLogReason("Velox anti-nuke: trusted role granted by someone untrusted")); err != nil {
		log.Printf("security: taking back trusted roles from %s in guild %s: %v", userID, guildID, err)
	}
}

// isBotRole reports whether roleID is one of Velox's own roles.
func (e *Engine) isBotRole(guildID, roleID string) bool {
	m, err := e.d.GuildMember(guildID, e.botID)
	if err != nil {
		return false
	}
	for _, r := range m.Roles {
		if r == roleID {
			return true
		}
	}
	return false
}