Catalogger.NET/Catalogger.Backend/Database/Repositories/MessageRepository.cs

// Copyright (C) 2021-present sam (starshines.gay)
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

using System.Text.Json;
using Catalogger.Backend.Extensions;
using Dapper;
using Remora.Discord.API;
using Remora.Discord.API.Abstractions.Gateway.Events;
using Remora.Discord.API.Abstractions.Objects;
using Remora.Rest.Core;

namespace Catalogger.Backend.Database.Repositories;

public class MessageRepository(
    ILogger logger,
    DatabaseConnection conn,
    IEncryptionService encryptionService
) : IDisposable, IAsyncDisposable
{
    private readonly ILogger _logger = logger.ForContext<MessageRepository>();

    public async Task<Message?> GetMessageAsync(ulong id, CancellationToken ct = default)
    {
        _logger.Debug("Retrieving message {MessageId}", id);

        var dbMsg = await conn.QueryFirstOrDefaultAsync<Models.Message>(
            "select * from messages where id = @Id",
            new { Id = id }
        );
        if (dbMsg == null)
            return null;

        return new Message(
            dbMsg.Id,
            dbMsg.OriginalId,
            dbMsg.UserId,
            dbMsg.ChannelId,
            dbMsg.GuildId,
            dbMsg.Member,
            dbMsg.System,
            Username: await Task.Run(() => encryptionService.Decrypt(dbMsg.Username), ct),
            Content: await Task.Run(() => encryptionService.Decrypt(dbMsg.Content), ct),
            Metadata: dbMsg.Metadata != null
                ? JsonSerializer.Deserialize<Metadata>(
                    await Task.Run(() => encryptionService.Decrypt(dbMsg.Metadata), ct)
                )
                : null,
            dbMsg.AttachmentSize
        );
    }

    /// <summary>
    /// Adds a new message. If the message is already in the database, updates the existing message instead.
    /// </summary>
    public async Task<bool> SaveMessageAsync(
        IMessage msg,
        Optional<Snowflake> guildId,
        CancellationToken ct = default
    )
    {
        var content = await Task.Run(
            () =>
                encryptionService.Encrypt(
                    string.IsNullOrWhiteSpace(msg.Content) ? "None" : msg.Content
                ),
            ct
        );
        var username = await Task.Run(() => encryptionService.Encrypt(msg.Author.Tag()), ct);
        var metadata = await Task.Run(
            () =>
                encryptionService.Encrypt(
                    JsonSerializer.Serialize(
                        new Metadata(
                            IsWebhook: msg.WebhookID.HasValue,
                            msg.Attachments.Select(a => new Attachment(
                                a.Filename,
                                a.Size,
                                a.ContentType.Value
                            ))
                        )
                    )
                ),
            ct
        );

        // MessageUpdateResponder wants to know whether the message already existed, so query this *before* inserting.
        var exists = await conn.ExecuteScalarAsync<bool>(
            "select exists(select id from messages where id = @Id)",
            new { Id = msg.ID.Value }
        );

        await conn.ExecuteAsync(
            """
            insert into messages (id, user_id, channel_id, guild_id, username, content, metadata, attachment_size)
            values (@Id, @UserId, @ChannelId, @GuildId, @Username, @Content, @Metadata, @AttachmentSize)
            on conflict (id) do update set username = @Username, content = @Content, metadata = @Metadata
            """,
            new
            {
                Id = msg.ID.Value,
                UserId = msg.Author.ID.Value,
                ChannelId = msg.ChannelID.Value,
                GuildId = guildId.IsDefined(out var guildIdValue)
                    ? guildIdValue.Value
                    : (ulong?)null,
                Content = content,
                Username = username,
                Metadata = metadata,
                AttachmentSize = msg.Attachments.Select(a => a.Size).Sum(),
            }
        );

        return exists;
    }

    public async Task<(bool IsStored, bool HasProxyInfo)> HasProxyInfoAsync(ulong id)
    {
        _logger.Debug("Checking if message {MessageId} has proxy information", id);

        var msg = await conn.QueryFirstOrDefaultAsync<(ulong Id, ulong OriginalId)>(
            "select id, original_id from messages where id = @Id",
            new { Id = id }
        );
        return (msg.Id != 0, msg.OriginalId != 0);
    }

    /// <summary>
    /// Updates a stored message with PluralKit information.
    /// </summary>
    /// <returns>True if the message exists and was updated, false if it doesn't exist.</returns>
    public async Task<bool> SetProxiedMessageDataAsync(
        ulong id,
        ulong originalId,
        ulong authorId,
        string? systemId,
        string? memberId
    )
    {
        _logger.Debug("Setting proxy information for message {MessageId}", id);

        var updatedCount = await conn.ExecuteAsync(
            "update messages set original_id = @OriginalId, user_id = @AuthorId, system = @SystemId, member = @MemberId where id = @Id",
            new
            {
                Id = id,
                OriginalId = originalId,
                AuthorId = authorId,
                SystemId = systemId,
                MemberId = memberId,
            }
        );
        if (updatedCount == 0)
        {
            _logger.Debug("Message {MessageId} not found, can't set proxy data for it", id);
            return false;
        }
        return true;
    }

    public async Task<bool> IsMessageIgnoredAsync(ulong id) =>
        await conn.ExecuteScalarAsync<bool>(
            "select exists(select id from ignored_messages where id = @Id)",
            new { Id = id }
        );

    public const int MaxMessageAgeDays = 15;

    public async Task<(int Messages, int IgnoredMessages)> DeleteExpiredMessagesAsync()
    {
        var cutoff = DateTimeOffset.UtcNow - TimeSpan.FromDays(MaxMessageAgeDays);
        var cutoffId = Snowflake.CreateTimestampSnowflake(cutoff, Constants.DiscordEpoch).Value;

        var msgCount = await conn.ExecuteAsync(
            "delete from messages where id < @Cutoff",
            new { Cutoff = cutoffId }
        );
        var ignoredMsgCount = await conn.ExecuteAsync(
            "delete from ignored_messages where id < @Cutoff",
            new { Cutoff = cutoffId }
        );

        return (msgCount, ignoredMsgCount);
    }

    public async Task IgnoreMessageAsync(ulong id) =>
        await conn.ExecuteAsync(
            "insert into ignored_messages (id) values (@Id) on conflict do nothing",
            new { Id = id }
        );

    public record Message(
        ulong Id,
        ulong? OriginalId,
        ulong UserId,
        ulong ChannelId,
        ulong GuildId,
        string? Member,
        string? System,
        string Username,
        string Content,
        Metadata? Metadata,
        int AttachmentSize
    );

    public record Metadata(bool IsWebhook, IEnumerable<Attachment> Attachments);

    public record Attachment(string Filename, int Size, string ContentType);

    public void Dispose()
    {
        conn.Dispose();
        GC.SuppressFinalize(this);
    }

    public async ValueTask DisposeAsync()
    {
        await conn.DisposeAsync();
        GC.SuppressFinalize(this);
    }
}