velox-bot/internal/security/restore.go

package security

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"sort"
	"time"

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

	"github.com/bwmarrin/discordgo"
)

// Snapshots record a guild's roles and channels (names, settings,
// permissions, channel permission overwrites) so a nuke can be undone.
// Messages and who had which role can't be captured or brought back; Discord
// offers no way to.
const (
	snapshotEvery    = 6 * time.Hour
	snapshotCheckGap = 30 * time.Minute
)

type snapRole struct {
	ID          string `json:"id"`
	Name        string `json:"name"`
	Color       int    `json:"color"`
	Hoist       bool   `json:"hoist"`
	Mentionable bool   `json:"mentionable"`
	Permissions int64  `json:"permissions"`
	Position    int    `json:"position"`
}

type snapOverwrite struct {
	ID    string `json:"id"`
	Type  int    `json:"type"`
	Allow int64  `json:"allow"`
	Deny  int64  `json:"deny"`
}

type snapChannel struct {
	ID         string          `json:"id"`
	Name       string          `json:"name"`
	Type       int             `json:"type"`
	Topic      string          `json:"topic,omitempty"`
	ParentID   string          `json:"parent_id,omitempty"`
	Position   int             `json:"position"`
	NSFW       bool            `json:"nsfw,omitempty"`
	RateLimit  int             `json:"rate_limit,omitempty"`
	Bitrate    int             `json:"bitrate,omitempty"`
	UserLimit  int             `json:"user_limit,omitempty"`
	Overwrites []snapOverwrite `json:"overwrites,omitempty"`
}

type snapshotData struct {
	Roles    []snapRole    `json:"roles"`
	Channels []snapChannel `json:"channels"`
}

// restorable channel types: no threads, which belong to their parent.
func restorable(t discordgo.ChannelType) bool {
	switch t {
	case discordgo.ChannelTypeGuildText, discordgo.ChannelTypeGuildVoice, discordgo.ChannelTypeGuildCategory,
		discordgo.ChannelTypeGuildNews, discordgo.ChannelTypeGuildStageVoice, discordgo.ChannelTypeGuildForum:
		return true
	}
	return false
}

// TakeSnapshot records the guild's current roles and channels.
func (e *Engine) TakeSnapshot(ctx context.Context, guildID, by string) (roles, channels int, err error) {
	rs, err := e.d.GuildRoles(guildID)
	if err != nil {
		return 0, 0, fmt.Errorf("listing roles: %w", err)
	}
	cs, err := e.d.GuildChannels(guildID)
	if err != nil {
		return 0, 0, fmt.Errorf("listing channels: %w", err)
	}
	var data snapshotData
	for _, r := range rs {
		if r.Managed {
			continue // integration roles come back with their bot
		}
		data.Roles = append(data.Roles, snapRole{ID: r.ID, Name: r.Name, Color: r.Color, Hoist: r.Hoist,
			Mentionable: r.Mentionable, Permissions: r.Permissions, Position: r.Position})
	}
	for _, c := range cs {
		if !restorable(c.Type) {
			continue
		}
		sc := snapChannel{ID: c.ID, Name: c.Name, Type: int(c.Type), Topic: c.Topic, ParentID: c.ParentID, Position: c.Position,
			NSFW: c.NSFW, RateLimit: c.RateLimitPerUser, Bitrate: c.Bitrate, UserLimit: c.UserLimit}
		for _, o := range c.PermissionOverwrites {
			sc.Overwrites = append(sc.Overwrites, snapOverwrite{ID: o.ID, Type: int(o.Type), Allow: o.Allow, Deny: o.Deny})
		}
		data.Channels = append(data.Channels, sc)
	}
	raw, err := json.Marshal(data)
	if err != nil {
		return 0, 0, err
	}
	if err := e.st.SaveSnapshot(ctx, securityrepo.Snapshot{GuildID: parseID(guildID), Data: raw, Roles: len(data.Roles), Channels: len(data.Channels)}); err != nil {
		return 0, 0, err
	}
	if by != "" {
		e.event(ctx, guildID, securityrepo.EventSnapshot, by, "", fmt.Sprintf("%d roles, %d channels", len(data.Roles), len(data.Channels)))
	}
	return len(data.Roles), len(data.Channels), nil
}

// WatchSnapshots takes a snapshot of every guild with anti-nuke on whose
// newest one is older than six hours. It runs until ctx ends.
func (e *Engine) WatchSnapshots(ctx context.Context) {
	run := func() {
		due, err := e.st.GuildsDueForSnapshot(ctx, snapshotEvery)
		if err != nil {
			log.Printf("security: finding guilds due a snapshot: %v", err)
			return
		}
		for _, g := range due {
			if _, _, err := e.TakeSnapshot(ctx, id(g), ""); err != nil {
				log.Printf("security: snapshot of guild %d: %v", g, err)
			}
		}
	}
	run()
	t := time.NewTicker(snapshotCheckGap)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			run()
		}
	}
}

// RestoreResult says what a restore did.
type RestoreResult struct {
	RolesCreated, RolesReverted, ChannelsCreated, Failed int
	SnapshotAt                                           time.Time
}

// Restore brings the guild back towards its latest snapshot:
//
//   - roles that are gone are recreated (same name, colour, permissions);
//   - roles that gained Administrator, Manage Server, Manage Roles or Manage
//     Webhooks since the snapshot get their snapshot permissions back;
//   - channels that are gone are recreated, categories first, with their
//     permission overwrites pointing at the recreated roles.
//
// Nothing is deleted: channels and roles made since the snapshot stay, since
// some may be legitimate. Restoring rewrites the server, so callers must check
// the actor is the owner or trusted. One restore per guild runs at a time.
func (e *Engine) Restore(ctx context.Context, guildID, by string) (*RestoreResult, error) {
	e.mu.Lock()
	if e.restoring[guildID] {
		e.mu.Unlock()
		return nil, UserError("A restore is already running for this server.")
	}
	e.restoring[guildID] = true
	e.mu.Unlock()
	defer func() {
		e.mu.Lock()
		delete(e.restoring, guildID)
		e.mu.Unlock()
	}()

	snap, err := e.st.LatestSnapshot(ctx, parseID(guildID))
	if err != nil {
		return nil, err
	}
	if snap == nil {
		return nil, UserError("There's no snapshot to restore from yet. Turn anti-nuke on (snapshots are taken every 6 hours) or take one with `/security snapshot`.")
	}
	var data snapshotData
	if err := json.Unmarshal(snap.Data, &data); err != nil {
		return nil, fmt.Errorf("reading the snapshot: %w", err)
	}
	res := &RestoreResult{SnapshotAt: snap.CreatedAt}
	reason := discordgo.WithAuditLogReason("Velox restore from snapshot")

	// Roles.
	current, err := e.d.GuildRoles(guildID)
	if err != nil {
		return nil, fmt.Errorf("listing roles: %w", err)
	}
	now := map[string]*discordgo.Role{}
	for _, r := range current {
		now[r.ID] = r
	}
	roleMap := map[string]string{} // snapshot role ID -> current role ID
	sort.Slice(data.Roles, func(i, j int) bool { return data.Roles[i].Position < data.Roles[j].Position })
	for _, sr := range data.Roles {
		if cur, ok := now[sr.ID]; ok {
			roleMap[sr.ID] = sr.ID
			if added := cur.Permissions &^ sr.Permissions; added&DangerousPerms != 0 && !cur.Managed {
				perms := sr.Permissions
				if _, err := e.d.GuildRoleEdit(guildID, sr.ID, &discordgo.RoleParams{Permissions: &perms}, reason); err != nil {
					log.Printf("security: restore: reverting role %s in guild %s: %v", sr.ID, guildID, err)
					res.Failed++
				} else {
					res.RolesReverted++
				}
			}
			continue
		}
		if sr.ID == guildID {
			continue // @everyone always exists
		}
		color, hoist, mention, perms := sr.Color, sr.Hoist, sr.Mentionable, sr.Permissions
		r, err := e.d.GuildRoleCreate(guildID, &discordgo.RoleParams{Name: sr.Name, Color: &color, Hoist: &hoist, Mentionable: &mention, Permissions: &perms}, reason)
		if err != nil {
			log.Printf("security: restore: recreating role %q in guild %s: %v", sr.Name, guildID, err)
			res.Failed++
			continue
		}
		roleMap[sr.ID] = r.ID
		res.RolesCreated++
	}

	// Channels: categories first so the others can be put back in them.
	chans, err := e.d.GuildChannels(guildID)
	if err != nil {
		return res, fmt.Errorf("listing channels: %w", err)
	}
	exists := map[string]bool{}
	for _, c := range chans {
		exists[c.ID] = true
	}
	chanMap := map[string]string{}
	for _, c := range chans {
		chanMap[c.ID] = c.ID
	}
	sort.SliceStable(data.Channels, func(i, j int) bool {
		ci := data.Channels[i].Type == int(discordgo.ChannelTypeGuildCategory)
		cj := data.Channels[j].Type == int(discordgo.ChannelTypeGuildCategory)
		if ci != cj {
			return ci
		}
		return data.Channels[i].Position < data.Channels[j].Position
	})
	for _, sc := range data.Channels {
		if exists[sc.ID] {
			continue
		}
		var overwrites []*discordgo.PermissionOverwrite
		for _, o := range sc.Overwrites {
			target := o.ID
			if discordgo.PermissionOverwriteType(o.Type) == discordgo.PermissionOverwriteTypeRole {
				mapped, ok := roleMap[o.ID]
				if !ok {
					continue // the role couldn't be brought back
				}
				target = mapped
			}
			overwrites = append(overwrites, &discordgo.PermissionOverwrite{ID: target, Type: discordgo.PermissionOverwriteType(o.Type), Allow: o.Allow, Deny: o.Deny})
		}
		parent := ""
		if sc.ParentID != "" {
			parent = chanMap[sc.ParentID] // "" if the category couldn't be restored
		}
		c, err := e.d.GuildChannelCreateComplex(guildID, discordgo.GuildChannelCreateData{
			Name: sc.Name, Type: discordgo.ChannelType(sc.Type), Topic: sc.Topic, Bitrate: sc.Bitrate, UserLimit: sc.UserLimit,
			RateLimitPerUser: sc.RateLimit, Position: sc.Position, PermissionOverwrites: overwrites, ParentID: parent, NSFW: sc.NSFW,
		}, reason)
		if err != nil {
			log.Printf("security: restore: recreating channel %q in guild %s: %v", sc.Name, guildID, err)
			res.Failed++
			continue
		}
		chanMap[sc.ID] = c.ID
		res.ChannelsCreated++
	}

	// Point the snapshot at the recreated roles and channels, so restoring
	// again doesn't create them a second time. The content stays the
	// snapshot's: nothing made since (or by an attacker) is added.
	if res.RolesCreated+res.ChannelsCreated > 0 {
		for i := range data.Roles {
			if nid, ok := roleMap[data.Roles[i].ID]; ok {
				data.Roles[i].ID = nid
			}
		}
		for i := range data.Channels {
			c := &data.Channels[i]
			if nid, ok := chanMap[c.ID]; ok {
				c.ID = nid
			}
			if nid, ok := chanMap[c.ParentID]; ok && c.ParentID != "" {
				c.ParentID = nid
			}
			for j := range c.Overwrites {
				if c.Overwrites[j].Type == int(discordgo.PermissionOverwriteTypeRole) {
					if nid, ok := roleMap[c.Overwrites[j].ID]; ok {
						c.Overwrites[j].ID = nid
					}
				}
			}
		}
		if raw, err := json.Marshal(data); err == nil {
			if err := e.st.SaveSnapshot(ctx, securityrepo.Snapshot{GuildID: parseID(guildID), Data: raw, Roles: snap.Roles, Channels: snap.Channels}); err != nil {
				log.Printf("security: saving the remapped snapshot for guild %s: %v", guildID, err)
			}
		}
	}

	e.event(ctx, guildID, securityrepo.EventRestored, by, "", fmt.Sprintf("%d roles recreated, %d reverted, %d channels recreated, %d failed",
		res.RolesCreated, res.RolesReverted, res.ChannelsCreated, res.Failed))
	if s, _, err := e.load(ctx, guildID); err == nil {
		e.send(ctx, guildID, s, alert{critical: true, embed: &discordgo.MessageEmbed{
			Title: "♻️ Server restored from snapshot",
			Description: fmt.Sprintf("Snapshot from <t:%d:f>. Recreated %d roles and %d channels, reverted dangerous permissions on %d roles. %d steps failed.\nMembers need their recreated roles given back by hand; Discord doesn't let bots see who had them.",
				res.SnapshotAt.Unix(), res.RolesCreated, res.ChannelsCreated, res.RolesReverted, res.Failed),
			Color: colorOK,
		}})
	}
	return res, nil
}