Catalogger.NET/go-exporter/message.go

package main

import (
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/diamondburned/arikawa/v3/discord"
	"github.com/georgysavva/scany/pgxscan"
)

func exportMessages() error {
	dirname := fmt.Sprintf("messages-%v.json", time.Now().Unix())
	err := os.Mkdir(dirname, 0o755)
	if err != nil {
		return fmt.Errorf("creating output directory: %w", err)
	}

	var ignoredMessages []int64
	err = conn.QueryRow(ctx, "select array_agg(id) from ignored_messages").Scan(&ignoredMessages)
	if err != nil {
		return fmt.Errorf("querying ignored messages: %w", err)
	}

	b, err := json.Marshal(ignoredMessages)
	if err != nil {
		return fmt.Errorf("marshaling ignored messages: %w", err)
	}

	err = os.WriteFile(dirname+"/ignored.json", b, 0o755)
	if err != nil {
		return fmt.Errorf("writing ignored messages: %w", err)
	}

	cursor := discord.Snowflake(0)
	for {
		start := time.Now()

		messages, err := getMessages(cursor)
		if err != nil {
			return fmt.Errorf("querying messages after %v: %w", cursor, err)
		}

		b, err := json.Marshal(messages)
		if err != nil {
			return fmt.Errorf("marshaling ignored messages: %w", err)
		}

		cursor = discord.Snowflake(messages[len(messages)-1].MsgID)

		filename := fmt.Sprintf("messages-%v-%v.json", messages[0].MsgID, cursor)

		err = os.WriteFile(dirname+"/"+filename, b, 0o755)
		if err != nil {
			return fmt.Errorf("writing messages: %w", err)
		}

		end := time.Now()

		log.Printf("Exported %v messages (starting at %v) in %v\n\n", len(messages), messages[0].MsgID, end.Sub(start))
		log.Printf("Current cursor: %v\n\n", cursor)

		if len(messages) < pageSize {
			break
		}
	}

	return nil
}

// Message is a single message
type Message struct {
	MsgID     discord.MessageID
	UserID    discord.UserID
	ChannelID discord.ChannelID
	ServerID  discord.GuildID

	Content  string
	Username string

	// These are only filled if the message was proxied by PluralKit
	Member *string
	System *string

	Metadata    *Metadata `db:"-"`
	RawMetadata *[]byte   `db:"metadata" json:"-"`
}

// Metadata is optional message metadata
type Metadata struct {
	UserID   *discord.UserID `json:"user_id,omitempty"`
	Username string          `json:"username,omitempty"`
	Avatar   string          `json:"avatar,omitempty"`
	Embeds   []discord.Embed `json:"embeds,omitempty"`
}

const pageSize = 10000

func getMessages(after discord.Snowflake) (ms []Message, err error) {
	err = pgxscan.Select(ctx, conn, &ms, "select * from messages where msg_id > $1 order by msg_id asc limit $2", after, pageSize)
	if err != nil {
		return nil, err
	}

	for i, m := range ms {
		b, err := hex.DecodeString(m.Content)
		if err != nil {
			return nil, err
		}

		out, err := Decrypt(b, aesKey)
		if err != nil {
			return nil, err
		}

		ms[i].Content = string(out)

		b, err = hex.DecodeString(m.Username)
		if err != nil {
			return nil, err
		}

		out, err = Decrypt(b, aesKey)
		if err != nil {
			return nil, err
		}

		ms[i].Username = string(out)

		if m.RawMetadata != nil {
			b, err := Decrypt(*m.RawMetadata, aesKey)
			if err == nil {
				var md Metadata
				err = json.Unmarshal(b, &md)
				if err != nil {
					return ms, err
				}
				ms[i].Metadata = &md
			}
		}
	}

	return ms, nil
}