velox-bot/internal/commands/security/public/security.go

package public

import (
	"context"
	"errors"
	"fmt"
	"log"
	"strings"
	"time"

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

	"github.com/bwmarrin/discordgo"
)

// Every handler re-checks permissions through the security engine. The
// DefaultMemberPermissions below only hide commands from the menu; server
// admins can change them, so they are never what keeps anyone out.

var (
	manageGuild   = int64(discordgo.PermissionManageGuild)
	manageChannel = int64(discordgo.PermissionManageChannels)
	guildOnly     = &[]discordgo.InteractionContextType{discordgo.InteractionContextGuild}
	textChannels  = []discordgo.ChannelType{discordgo.ChannelTypeGuildText, discordgo.ChannelTypeGuildNews}
	chatChannels  = []discordgo.ChannelType{
		discordgo.ChannelTypeGuildText, discordgo.ChannelTypeGuildNews, discordgo.ChannelTypeGuildForum,
		discordgo.ChannelTypeGuildVoice, discordgo.ChannelTypeGuildStageVoice,
	}
	opTimeout     = 3 * time.Minute
	genericFailed = "Something went wrong. Check that Velox's role is high enough and has the permissions it needs, then try again."
)

var featureChoices = []*discordgo.ApplicationCommandOptionChoice{
	{Name: "Join Gate", Value: securityrepo.FeatureJoinGate},
	{Name: "Raid mode", Value: securityrepo.FeatureAntiRaid},
	{Name: "Anti-nuke", Value: securityrepo.FeatureAntiNuke},
	{Name: "Heat (anti-spam)", Value: securityrepo.FeatureHeat},
}

var Security = &discordgo.ApplicationCommand{
	Name:                     "security",
	Description:              "Anti-raid and anti-nuke protection",
	DefaultMemberPermissions: &manageGuild,
	Contexts:                 guildOnly,
	Options: []*discordgo.ApplicationCommandOption{
		{Type: discordgo.ApplicationCommandOptionSubCommand, Name: "status", Description: "Show what's on and who is trusted"},
		{
			Type: discordgo.ApplicationCommandOptionSubCommand, Name: "toggle", Description: "Turn a protection on or off",
			Options: []*discordgo.ApplicationCommandOption{
				{Type: discordgo.ApplicationCommandOptionString, Name: "feature", Description: "Which protection", Required: true, Choices: featureChoices},
				{Type: discordgo.ApplicationCommandOptionBoolean, Name: "on", Description: "On or off", Required: true},
			},
		},
		{
			Type: discordgo.ApplicationCommandOptionSubCommand, Name: "alerts", Description: "Where alerts go and which role they ping (owner or trusted)",
			Options: []*discordgo.ApplicationCommandOption{
				{Type: discordgo.ApplicationCommandOptionChannel, Name: "channel", Description: "Alert channel", Required: true, ChannelTypes: textChannels},
				{Type: discordgo.ApplicationCommandOptionRole, Name: "role", Description: "Role to ping for raids and nukes"},
			},
		},
		{
			Type: discordgo.ApplicationCommandOptionSubCommand, Name: "quarantine-role", Description: "Role given to quarantined members (owner or trusted)",
			Options: []*discordgo.ApplicationCommandOption{
				{Type: discordgo.ApplicationCommandOptionRole, Name: "role", Description: "Leave empty to use none"},
			},
		},
		{
			Type: discordgo.ApplicationCommandOptionSubCommandGroup, Name: "trust", Description: "Who is exempt from anti-nuke (owner only)",
			Options: []*discordgo.ApplicationCommandOption{
				{
					Type: discordgo.ApplicationCommandOptionSubCommand, Name: "add", Description: "Trust a user or a role",
					Options: []*discordgo.ApplicationCommandOption{
						{Type: discordgo.ApplicationCommandOptionUser, Name: "user", Description: "A user"},
						{Type: discordgo.ApplicationCommandOptionRole, Name: "role", Description: "A role"},
					},
				},
				{
					Type: discordgo.ApplicationCommandOptionSubCommand, Name: "remove", Description: "Stop trusting a user or a role",
					Options: []*discordgo.ApplicationCommandOption{
						{Type: discordgo.ApplicationCommandOptionUser, Name: "user", Description: "A user"},
						{Type: discordgo.ApplicationCommandOptionRole, Name: "role", Description: "A role"},
					},
				},
			},
		},
		{
			Type: discordgo.ApplicationCommandOptionSubCommand, Name: "release", Description: "Release a quarantined member (owner or trusted)",
			Options: []*discordgo.ApplicationCommandOption{
				{Type: discordgo.ApplicationCommandOptionUser, Name: "user", Description: "Who to release", Required: true},
			},
		},
		{Type: discordgo.ApplicationCommandOptionSubCommand, Name: "quarantined", Description: "List quarantined members"},
		{Type: discordgo.ApplicationCommandOptionSubCommand, Name: "end-raid", Description: "End raid mode now (owner or trusted)"},
		{Type: discordgo.ApplicationCommandOptionSubCommand, Name: "snapshot", Description: "Save a snapshot of roles and channels to restore from"},
		{Type: discordgo.ApplicationCommandOptionSubCommand, Name: "restore", Description: "Recreate deleted roles and channels from the latest snapshot (owner or trusted)"},
	},
}

func lockOptions(unlock bool) []*discordgo.ApplicationCommandOption {
	verb := "Lock"
	if unlock {
		verb = "Unlock"
	}
	joins := &discordgo.ApplicationCommandOption{Type: discordgo.ApplicationCommandOptionSubCommand, Name: "joins", Description: verb + " joins"}
	if !unlock {
		joins.Description = "Kick or ban everyone who joins until unlocked"
		joins.Options = []*discordgo.ApplicationCommandOption{{
			Type: discordgo.ApplicationCommandOptionString, Name: "action", Description: "What happens to new members", Required: true,
			Choices: []*discordgo.ApplicationCommandOptionChoice{{Name: "Kick", Value: securityrepo.ActionKick}, {Name: "Ban", Value: securityrepo.ActionBan}},
		}}
	}
	return []*discordgo.ApplicationCommandOption{
		{
			Type: discordgo.ApplicationCommandOptionSubCommand, Name: "channel", Description: verb + " one channel",
			Options: []*discordgo.ApplicationCommandOption{{Type: discordgo.ApplicationCommandOptionChannel, Name: "channel", Description: "Defaults to this channel", ChannelTypes: chatChannels}},
		},
		{Type: discordgo.ApplicationCommandOptionSubCommand, Name: "server", Description: verb + " every channel"},
		joins,
	}
}

var Lockdown = &discordgo.ApplicationCommand{
	Name: "lockdown", Description: "Stop members from talking, or from joining", DefaultMemberPermissions: &manageChannel,
	Contexts: guildOnly, Options: lockOptions(false),
}

var Unlock = &discordgo.ApplicationCommand{
	Name: "unlock", Description: "Undo a lockdown (owner or trusted)", DefaultMemberPermissions: &manageChannel,
	Contexts: guildOnly, Options: lockOptions(true),
}

func reply(s *discordgo.Session, i *discordgo.InteractionCreate, msg string) {
	_ = s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
		Type: discordgo.InteractionResponseChannelMessageWithSource,
		Data: &discordgo.InteractionResponseData{Content: msg, Flags: discordgo.MessageFlagsEphemeral, AllowedMentions: &discordgo.MessageAllowedMentions{}},
	})
}

func deferReply(s *discordgo.Session, i *discordgo.InteractionCreate) bool {
	return s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
		Type: discordgo.InteractionResponseDeferredChannelMessageWithSource,
		Data: &discordgo.InteractionResponseData{Flags: discordgo.MessageFlagsEphemeral},
	}) == nil
}

func edit(s *discordgo.Session, i *discordgo.InteractionCreate, msg string) {
	_, _ = s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{Content: &msg, AllowedMentions: &discordgo.MessageAllowedMentions{}})
}

func message(where string, err error) string {
	var ue security.UserError
	if errors.As(err, &ue) {
		return string(ue)
	}
	log.Printf("security command %s: %v", where, err)
	return genericFailed
}

func opts(o []*discordgo.ApplicationCommandInteractionDataOption) map[string]*discordgo.ApplicationCommandInteractionDataOption {
	m := map[string]*discordgo.ApplicationCommandInteractionDataOption{}
	for _, x := range o {
		m[x.Name] = x
	}
	return m
}

func perms(i *discordgo.InteractionCreate) int64 {
	if i.Member == nil {
		return 0
	}
	return i.Member.Permissions
}

func userID(i *discordgo.InteractionCreate) string {
	if i.Member == nil || i.Member.User == nil {
		return ""
	}
	return i.Member.User.ID
}

func engine(s *discordgo.Session, i *discordgo.InteractionCreate) *security.Engine {
	e := security.Default()
	if e == nil || i.GuildID == "" || i.Member == nil {
		reply(s, i, "Security isn't available right now.")
		return nil
	}
	return e
}

// allowedOrTrusted lets an action through when the member has one of the
// permissions, or is the owner or trusted.
func allowedOrTrusted(ctx context.Context, e *security.Engine, i *discordgo.InteractionCreate, need int64) bool {
	if perms(i)&(need|discordgo.PermissionAdministrator) != 0 {
		return true
	}
	return e.RequireTrusted(ctx, i.GuildID, i.Member) == nil
}

func SecurityHandler(s *discordgo.Session, i *discordgo.InteractionCreate) {
	e := engine(s, i)
	if e == nil {
		return
	}
	data := i.ApplicationCommandData()
	if len(data.Options) == 0 {
		return
	}
	sub := data.Options[0]
	args := opts(sub.Options)
	ctx, cancel := context.WithTimeout(context.Background(), opTimeout)
	defer cancel()

	switch sub.Name {
	case "status":
		if !allowedOrTrusted(ctx, e, i, discordgo.PermissionManageGuild) {
			reply(s, i, "You need Manage Server to see the security settings.")
			return
		}
		reply(s, i, status(ctx, e, i.GuildID))
	case "toggle":
		feature, on := args["feature"].StringValue(), args["on"].BoolValue()
		if err := e.SetFeature(ctx, i.GuildID, i.Member, perms(i), feature, on); err != nil {
			reply(s, i, message("toggle", err))
			return
		}
		state := "off"
		if on {
			state = "on"
		}
		reply(s, i, fmt.Sprintf("%s is now **%s**.", featureName(feature), state))
	case "alerts":
		ch := args["channel"].ChannelValue(nil).ID
		role := ""
		if r, ok := args["role"]; ok {
			role = r.RoleValue(nil, i.GuildID).ID
		}
		if err := e.SetAlerts(ctx, i.GuildID, i.Member, ch, role); err != nil {
			reply(s, i, message("alerts", err))
			return
		}
		reply(s, i, fmt.Sprintf("Alerts will go to <#%s>.", ch))
	case "quarantine-role":
		role := ""
		if r, ok := args["role"]; ok {
			role = r.RoleValue(nil, i.GuildID).ID
		}
		if err := e.SetQuarantineRole(ctx, i.GuildID, i.Member, role); err != nil {
			reply(s, i, message("quarantine-role", err))
			return
		}
		if role == "" {
			reply(s, i, "Quarantined members won't get a role; they still lose their roles and are timed out.")
		} else {
			reply(s, i, fmt.Sprintf("Quarantined members will get <@&%s>.", role))
		}
	case "trust":
		if len(sub.Options) == 0 {
			return
		}
		handleTrust(ctx, s, i, e, sub.Options[0])
	case "release":
		if err := e.RequireTrusted(ctx, i.GuildID, i.Member); err != nil {
			reply(s, i, message("release", err))
			return
		}
		target := args["user"].UserValue(nil)
		if !deferReply(s, i) {
			return
		}
		if err := e.Release(ctx, i.GuildID, target.ID, userID(i)); err != nil {
			edit(s, i, message("release", err))
			return
		}
		edit(s, i, fmt.Sprintf("<@%s> was released and got their roles back.", target.ID))
	case "quarantined":
		if !allowedOrTrusted(ctx, e, i, discordgo.PermissionManageGuild) {
			reply(s, i, "You need Manage Server to see this.")
			return
		}
		qs, err := e.Quarantines(ctx, i.GuildID)
		if err != nil {
			reply(s, i, message("quarantined", err))
			return
		}
		if len(qs) == 0 {
			reply(s, i, "Nobody is quarantined.")
			return
		}
		var b strings.Builder
		for _, q := range qs {
			fmt.Fprintf(&b, "• <@%d>, <t:%d:R>: %s\n", q.UserID, q.CreatedAt.Unix(), q.Reason)
		}
		reply(s, i, clip(b.String(), 1900))
	case "end-raid":
		endRaid(ctx, s, i, e)
	case "snapshot":
		if !allowedOrTrusted(ctx, e, i, discordgo.PermissionManageGuild) {
			reply(s, i, "You need Manage Server to take a snapshot.")
			return
		}
		if !deferReply(s, i) {
			return
		}
		roles, channels, err := e.TakeSnapshot(ctx, i.GuildID, userID(i))
		if err != nil {
			edit(s, i, message("snapshot", err))
			return
		}
		edit(s, i, fmt.Sprintf("Snapshot saved: %d roles and %d channels. The last three are kept.", roles, channels))
	case "restore":
		if err := e.RequireTrusted(ctx, i.GuildID, i.Member); err != nil {
			reply(s, i, message("restore", err))
			return
		}
		if !deferReply(s, i) {
			return
		}
		// Recreating many channels takes a while under Discord's rate limits.
		res, err := e.Restore(ctx, i.GuildID, userID(i))
		if err != nil {
			edit(s, i, message("restore", err))
			return
		}
		edit(s, i, fmt.Sprintf("Restored from the snapshot of <t:%d:f>: %d roles and %d channels recreated, %d roles had dangerous permissions reverted, %d steps failed.",
			res.SnapshotAt.Unix(), res.RolesCreated, res.ChannelsCreated, res.RolesReverted, res.Failed))
	}
}

func handleTrust(ctx context.Context, s *discordgo.Session, i *discordgo.InteractionCreate, e *security.Engine, sub *discordgo.ApplicationCommandInteractionDataOption) {
	args := opts(sub.Options)
	kind, target := "", ""
	if u, ok := args["user"]; ok {
		kind, target = "user", u.UserValue(nil).ID
	}
	if r, ok := args["role"]; ok {
		if target != "" {
			reply(s, i, "Pick a user or a role, not both.")
			return
		}
		kind, target = "role", r.RoleValue(nil, i.GuildID).ID
	}
	if target == "" {
		reply(s, i, "Pick a user or a role.")
		return
	}
	mention := "<@" + target + ">"
	if kind == "role" {
		mention = "<@&" + target + ">"
	}
	switch sub.Name {
	case "add":
		if err := e.Trust(ctx, i.GuildID, i.Member, kind, target); err != nil {
			reply(s, i, message("trust add", err))
			return
		}
		reply(s, i, mention+" is now trusted: exempt from anti-nuke and able to change security settings.")
	case "remove":
		if err := e.Untrust(ctx, i.GuildID, i.Member, target); err != nil {
			reply(s, i, message("trust remove", err))
			return
		}
		reply(s, i, mention+" is no longer trusted.")
	}
}

func endRaid(ctx context.Context, s *discordgo.Session, i *discordgo.InteractionCreate, e *security.Engine) {
	if err := e.RequireTrusted(ctx, i.GuildID, i.Member); err != nil {
		reply(s, i, message("end raid", err))
		return
	}
	if !deferReply(s, i) {
		return
	}
	if err := e.EndRaid(ctx, i.GuildID, userID(i)); err != nil {
		edit(s, i, message("end raid", err))
		return
	}
	edit(s, i, "Raid mode is off.")
}

func featureName(f string) string {
	switch f {
	case securityrepo.FeatureJoinGate:
		return "Join Gate"
	case securityrepo.FeatureAntiRaid:
		return "Raid mode"
	case securityrepo.FeatureAntiNuke:
		return "Anti-nuke"
	case securityrepo.FeatureHeat:
		return "Heat"
	}
	return f
}

func onOff(b bool) string {
	if b {
		return "✅ on"
	}
	return "⬜ off"
}

func status(ctx context.Context, e *security.Engine, guildID string) string {
	st, trusted, err := e.Settings(ctx, guildID)
	if err != nil {
		return message("status", err)
	}
	var b strings.Builder
	fmt.Fprintf(&b, "**Join Gate:** %s\n**Raid mode:** %s (%d joins in %ds → %s)\n**Anti-nuke:** %s (auto restore %s)\n**Heat:** %s (%s)\n",
		onOff(st.JoinGateEnabled), onOff(st.AntiRaidEnabled), st.RaidJoinCount, st.RaidJoinSeconds, st.RaidAction,
		onOff(st.AntiNukeEnabled), onOff(st.AntiNukeAutoRestore), onOff(st.HeatEnabled), st.HeatStrictness)
	if e.RaidActive(guildID) {
		b.WriteString("🚨 **A raid is in progress.**\n")
	}
	if st.JoinLock != securityrepo.ActionOff {
		fmt.Fprintf(&b, "🔒 Joins are locked (%s).\n", st.JoinLock)
	}
	if st.AlertChannelID != 0 {
		fmt.Fprintf(&b, "**Alerts:** <#%d>", st.AlertChannelID)
		if st.AlertRoleID != 0 {
			fmt.Fprintf(&b, ", pinging <@&%d>", st.AlertRoleID)
		}
		b.WriteString("\n")
	} else {
		b.WriteString("**Alerts:** not set. Use `/security alerts`.\n")
	}
	if owner, err := e.OwnerID(guildID); err == nil {
		fmt.Fprintf(&b, "**Trusted:** <@%s> (owner)", owner)
	}
	for _, t := range trusted {
		if t.Kind == "role" {
			fmt.Fprintf(&b, ", <@&%d>", t.TargetID)
		} else {
			fmt.Fprintf(&b, ", <@%d>", t.TargetID)
		}
	}
	b.WriteString("\nFine-tune filters, limits and actions on the dashboard.")
	return b.String()
}

func clip(s string, n int) string {
	if r := []rune(s); len(r) > n {
		return string(r[:n-1]) + "…"
	}
	return s
}

func LockdownHandler(s *discordgo.Session, i *discordgo.InteractionCreate) {
	e := engine(s, i)
	if e == nil {
		return
	}
	data := i.ApplicationCommandData()
	if len(data.Options) == 0 {
		return
	}
	sub := data.Options[0]
	args := opts(sub.Options)
	ctx, cancel := context.WithTimeout(context.Background(), opTimeout)
	defer cancel()

	switch sub.Name {
	case "channel":
		if !allowedOrTrusted(ctx, e, i, discordgo.PermissionManageChannels) {
			reply(s, i, "You need Manage Channels to lock a channel.")
			return
		}
		ch := i.ChannelID
		if c, ok := args["channel"]; ok {
			ch = c.ChannelValue(nil).ID
		}
		if err := e.LockChannel(ctx, i.GuildID, ch, userID(i), false); err != nil {
			reply(s, i, message("lock channel", err))
			return
		}
		reply(s, i, fmt.Sprintf("🔒 <#%s> is locked. `/unlock channel` puts it back as it was.", ch))
	case "server":
		if !allowedOrTrusted(ctx, e, i, discordgo.PermissionManageGuild) {
			reply(s, i, "You need Manage Server to lock the server.")
			return
		}
		if !deferReply(s, i) {
			return
		}
		locked, failed, err := e.LockServer(ctx, i.GuildID, userID(i))
		if err != nil {
			edit(s, i, message("lock server", err))
			return
		}
		msg := fmt.Sprintf("🔒 Locked %d channels.", locked)
		if failed > 0 {
			msg += fmt.Sprintf(" %d couldn't be locked (Velox needs Manage Channels there).", failed)
		}
		edit(s, i, msg+" `/unlock server` puts them back as they were.")
	case "joins":
		if !allowedOrTrusted(ctx, e, i, discordgo.PermissionManageGuild) {
			reply(s, i, "You need Manage Server to lock joins.")
			return
		}
		mode := args["action"].StringValue()
		if err := e.SetJoinLock(ctx, i.GuildID, mode, userID(i)); err != nil {
			reply(s, i, message("lock joins", err))
			return
		}
		reply(s, i, fmt.Sprintf("🔒 Everyone who joins will be %s until `/unlock joins`.", map[string]string{"kick": "kicked", "ban": "banned"}[mode]))
	}
}

// UnlockHandler relaxes protection, so it is owner or trusted only.
func UnlockHandler(s *discordgo.Session, i *discordgo.InteractionCreate) {
	e := engine(s, i)
	if e == nil {
		return
	}
	data := i.ApplicationCommandData()
	if len(data.Options) == 0 {
		return
	}
	sub := data.Options[0]
	args := opts(sub.Options)
	ctx, cancel := context.WithTimeout(context.Background(), opTimeout)
	defer cancel()
	if err := e.RequireTrusted(ctx, i.GuildID, i.Member); err != nil {
		reply(s, i, message("unlock", err))
		return
	}

	switch sub.Name {
	case "channel":
		ch := i.ChannelID
		if c, ok := args["channel"]; ok {
			ch = c.ChannelValue(nil).ID
		}
		if err := e.UnlockChannel(ctx, i.GuildID, ch, userID(i)); err != nil {
			reply(s, i, message("unlock channel", err))
			return
		}
		reply(s, i, fmt.Sprintf("🔓 <#%s> is unlocked.", ch))
	case "server":
		if !deferReply(s, i) {
			return
		}
		unlocked, failed, err := e.UnlockServer(ctx, i.GuildID, userID(i))
		if err != nil {
			edit(s, i, message("unlock server", err))
			return
		}
		msg := fmt.Sprintf("🔓 Unlocked %d channels.", unlocked)
		if failed > 0 {
			msg += fmt.Sprintf(" %d failed; try `/unlock server` again.", failed)
		}
		edit(s, i, msg)
	case "joins":
		if err := e.SetJoinLock(ctx, i.GuildID, securityrepo.ActionOff, userID(i)); err != nil {
			reply(s, i, message("unlock joins", err))
			return
		}
		reply(s, i, "🔓 Joins are open again.")
	}
}

// HandleComponent handles the alert's "End raid" button. Anyone in the alert
// channel can press it, so it checks trust like the command.
func HandleComponent(s *discordgo.Session, i *discordgo.InteractionCreate) {
	e := engine(s, i)
	if e == nil {
		return
	}
	if i.MessageComponentData().CustomID != security.CustomIDEndRaid {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), opTimeout)
	defer cancel()
	endRaid(ctx, s, i, e)
}