Catalogger.NET/Catalogger.Backend/Database/Queries/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 Microsoft.EntityFrameworkCore;
using Remora.Discord.API;
using Remora.Discord.API.Abstractions.Gateway.Events;
using Remora.Rest.Core;
using DbMessage = Catalogger.Backend.Database.Models.Message;

namespace Catalogger.Backend.Database.Queries;

public class MessageRepository(
    ILogger logger,
    DatabaseContext db,
    IEncryptionService encryptionService
)
{
    private readonly ILogger _logger = logger.ForContext<MessageRepository>();

    public async Task SaveMessageAsync(IMessageCreate msg, CancellationToken ct = default)
    {
        _logger.Debug("Saving message {MessageId}", msg.ID);

        var metadata = new Metadata(
            IsWebhook: msg.WebhookID.HasValue,
            msg.Attachments.Select(a => new Attachment(a.Filename, a.Size, a.ContentType.Value))
        );

        var dbMessage = new DbMessage
        {
            Id = msg.ID.ToUlong(),
            UserId = msg.Author.ID.ToUlong(),
            ChannelId = msg.ChannelID.ToUlong(),
            GuildId = msg.GuildID.ToUlong(),

            Content = await Task.Run(
                () =>
                    encryptionService.Encrypt(
                        string.IsNullOrWhiteSpace(msg.Content) ? "None" : msg.Content
                    ),
                ct
            ),
            Username = await Task.Run(() => encryptionService.Encrypt(msg.Author.Tag()), ct),
            Metadata = await Task.Run(
                () => encryptionService.Encrypt(JsonSerializer.Serialize(metadata)),
                ct
            ),
            AttachmentSize = msg.Attachments.Select(a => a.Size).Sum(),
        };

        db.Add(dbMessage);
        await db.SaveChangesAsync(ct);
    }

    /// <summary>
    /// Updates an edited message.
    /// </summary>
    /// <returns>true if the message was already stored and got updated,
    /// false if the message wasn't stored and was newly inserted.</returns>
    public async Task<bool> UpdateMessageAsync(IMessageCreate msg, CancellationToken ct = default)
    {
        _logger.Debug("Updating message {MessageId}", msg.ID);

        var tx = await db.Database.BeginTransactionAsync(ct);
        var (isStored, _) = await HasProxyInfoAsync(msg.ID.Value);
        if (!isStored)
        {
            _logger.Debug("Edited message {MessageId} is not stored yet, storing it", msg.ID);
            await SaveMessageAsync(msg, ct);
            await tx.CommitAsync(ct);
            return false;
        }
        else
        {
            var metadata = new Metadata(
                IsWebhook: msg.WebhookID.HasValue,
                msg.Attachments.Select(a => new Attachment(a.Filename, a.Size, a.ContentType.Value))
            );

            var dbMsg = await db.Messages.FindAsync(
                new object?[] { msg.ID.Value },
                cancellationToken: ct
            );
            if (dbMsg == null)
                throw new CataloggerError(
                    "Message was null despite HasProxyInfoAsync returning true"
                );

            dbMsg.Content = await Task.Run(
                () =>
                    encryptionService.Encrypt(
                        string.IsNullOrWhiteSpace(msg.Content) ? "None" : msg.Content
                    ),
                ct
            );
            dbMsg.Username = await Task.Run(() => encryptionService.Encrypt(msg.Author.Tag()), ct);
            dbMsg.Metadata = await Task.Run(
                () => encryptionService.Encrypt(JsonSerializer.Serialize(metadata)),
                ct
            );

            db.Update(dbMsg);
            await db.SaveChangesAsync(ct);
            await tx.CommitAsync(ct);
            return true;
        }
    }

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

        var dbMsg = await db.Messages.AsNoTracking().FirstOrDefaultAsync(m => m.Id == id, ct);
        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>
    /// Checks if a message has proxy information.
    /// If yes, returns (true, true). If no, returns (true, false). If the message isn't saved at all, returns (false, false).
    /// </summary>
    public async Task<(bool, bool)> HasProxyInfoAsync(ulong id)
    {
        _logger.Debug("Checking if message {MessageId} has proxy information", id);

        var msg = await db
            .Messages.AsNoTracking()
            .Select(m => new { m.Id, m.OriginalId })
            .FirstOrDefaultAsync(m => m.Id == id);
        return (msg != null, msg?.OriginalId != null);
    }

    public async Task SetProxiedMessageDataAsync(
        ulong id,
        ulong originalId,
        ulong authorId,
        string? systemId,
        string? memberId
    )
    {
        _logger.Debug("Setting proxy information for message {MessageId}", id);

        var message = await db.Messages.FirstOrDefaultAsync(m => m.Id == id);
        if (message == null)
        {
            _logger.Debug("Message {MessageId} not found", id);
            return;
        }

        _logger.Debug("Updating message {MessageId}", id);

        message.OriginalId = originalId;
        message.UserId = authorId;
        message.System = systemId;
        message.Member = memberId;

        db.Update(message);
        await db.SaveChangesAsync();
    }

    public async Task<bool> IsMessageIgnoredAsync(ulong id, CancellationToken ct = default)
    {
        _logger.Debug("Checking if message {MessageId} is ignored", id);
        return await db.IgnoredMessages.AsNoTracking().FirstOrDefaultAsync(m => m.Id == id, ct)
            != null;
    }

    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);

        var msgCount = await db.Messages.Where(m => m.Id < cutoffId.Value).ExecuteDeleteAsync();
        var ignoredMsgCount = await db
            .IgnoredMessages.Where(m => m.Id < cutoffId.Value)
            .ExecuteDeleteAsync();

        return (msgCount, ignoredMsgCount);
    }

    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);
}