Cleanuparr/code/backend/Cleanuparr.Infrastructure/Features/Jobs/MalwareBlocker.cs

using Cleanuparr.Domain.Entities.Arr.Queue;
using Cleanuparr.Domain.Enums;
using Cleanuparr.Infrastructure.Events.Interfaces;
using Cleanuparr.Infrastructure.Features.Arr.Interfaces;
using Cleanuparr.Infrastructure.Features.Context;
using Cleanuparr.Infrastructure.Features.DownloadClient;
using Cleanuparr.Infrastructure.Features.MalwareBlocker;
using Cleanuparr.Infrastructure.Helpers;
using Cleanuparr.Infrastructure.Services.Interfaces;
using Cleanuparr.Persistence;
using Cleanuparr.Persistence.Models.Configuration;
using Cleanuparr.Persistence.Models.Configuration.Arr;
using Cleanuparr.Persistence.Models.Configuration.General;
using Cleanuparr.Persistence.Models.Configuration.MalwareBlocker;
using MassTransit;
using Microsoft.Extensions.Caching.Memory;
using Microsoft.Extensions.Logging;
using LogContext = Serilog.Context.LogContext;

namespace Cleanuparr.Infrastructure.Features.Jobs;

public sealed class MalwareBlocker : GenericHandler
{
    private readonly IBlocklistProvider _blocklistProvider;
    private readonly IJobManagementService _jobManagementService;

    public MalwareBlocker(
        ILogger<MalwareBlocker> logger,
        DataContext dataContext,
        IMemoryCache cache,
        IBus messageBus,
        IArrClientFactory arrClientFactory,
        IArrQueueIterator arrArrQueueIterator,
        IDownloadServiceFactory downloadServiceFactory,
        IBlocklistProvider blocklistProvider,
        IEventPublisher eventPublisher,
        IJobManagementService jobManagementService
    ) : base(
        logger, dataContext, cache, messageBus,
        arrClientFactory, arrArrQueueIterator, downloadServiceFactory, eventPublisher
    )
    {
        _blocklistProvider = blocklistProvider;
        _jobManagementService = jobManagementService;
    }

    protected override async Task ExecuteInternalAsync(CancellationToken cancellationToken = default)
    {
        if (ContextProvider.Get<List<DownloadClientConfig>>(nameof(DownloadClientConfig)).Count is 0)
        {
            _logger.LogWarning("No download clients configured");
            return;
        }
        
        ContentBlockerConfig malwareBlockerConfig = ContextProvider.Get<ContentBlockerConfig>();
        
        if (!malwareBlockerConfig.Sonarr.Enabled && 
            !malwareBlockerConfig.Radarr.Enabled && 
            !malwareBlockerConfig.Lidarr.Enabled && 
            !malwareBlockerConfig.Readarr.Enabled && 
            !malwareBlockerConfig.Whisparr.Enabled)
        {
            _logger.LogWarning("No blocklists are enabled");
            return;
        }

        await _blocklistProvider.LoadBlocklistsAsync();

        if (ContextProvider.Get(nameof(WebhookScanTarget)) is WebhookScanTarget webhookTarget)
        {
            await ProcessWebhookTargetAsync(malwareBlockerConfig, webhookTarget);
            return;
        }

        var sonarrConfig = ContextProvider.Get<ArrConfig>(nameof(InstanceType.Sonarr));
        var radarrConfig = ContextProvider.Get<ArrConfig>(nameof(InstanceType.Radarr));
        var lidarrConfig = ContextProvider.Get<ArrConfig>(nameof(InstanceType.Lidarr));
        var readarrConfig = ContextProvider.Get<ArrConfig>(nameof(InstanceType.Readarr));
        var whisparrConfig = ContextProvider.Get<ArrConfig>(nameof(InstanceType.Whisparr));

        if (malwareBlockerConfig.Sonarr.Enabled)
        {
            await ProcessArrConfigAsync(sonarrConfig);
        }

        if (malwareBlockerConfig.Radarr.Enabled)
        {
            await ProcessArrConfigAsync(radarrConfig);
        }

        if (malwareBlockerConfig.Lidarr.Enabled)
        {
            await ProcessArrConfigAsync(lidarrConfig);
        }

        if (malwareBlockerConfig.Readarr.Enabled)
        {
            await ProcessArrConfigAsync(readarrConfig);
        }

        if (malwareBlockerConfig.Whisparr.Enabled)
        {
            await ProcessArrConfigAsync(whisparrConfig);
        }
    }

    /// <summary>
    /// Scans a single download identified by an *arr "On Grab" webhook, restricted to the originating instance, instead of iterating the whole queue.
    /// Schedules the next retry only when the download was not yet found/scanned.
    /// </summary>
    private async Task ProcessWebhookTargetAsync(ContentBlockerConfig config, WebhookScanTarget target)
    {
        BlocklistSettings? blocklist = target.Type switch
        {
            InstanceType.Sonarr => config.Sonarr,
            InstanceType.Radarr => config.Radarr,
            _ => null,
        };

        if (blocklist is null || !blocklist.Enabled)
        {
            _logger.LogDebug("skip webhook scan | blocklist for {type} is not enabled", target.Type);
            return;
        }

        ArrConfig arrConfig = ContextProvider.Get<ArrConfig>(target.Type.ToString());
        ArrInstance? instance = arrConfig.Instances
            .FirstOrDefault(x => x.Id == target.InstanceId && x.Enabled);

        if (instance is null)
        {
            _logger.LogWarning("skip webhook scan | instance {id} not found or disabled", target.InstanceId);
            return;
        }

        instance.ArrConfig = arrConfig;

        bool resolved = await ScanInstanceAsync(instance, target);

        if (!resolved)
        {
            await _jobManagementService.ScheduleMalwareBlockerWebhookRetry(target);
        }
    }

    protected override Task ProcessInstanceAsync(ArrInstance instance) => ScanInstanceAsync(instance);

    private async Task<bool> ScanInstanceAsync(ArrInstance instance, WebhookScanTarget? target = null)
    {
        List<string> ignoredDownloads = ContextProvider.Get<GeneralConfig>(nameof(GeneralConfig)).IgnoredDownloads
            .Concat(ContextProvider.Get<ContentBlockerConfig>().IgnoredDownloads)
            .ToList();

        using var _ = LogContext.PushProperty(LogProperties.Category, instance.ArrConfig.Type.ToString());
        using var _2 = LogContext.PushProperty(LogProperties.InstanceName, instance.Name);

        IArrClient arrClient = _arrClientFactory.GetClient(instance.ArrConfig.Type, instance.Version);

        // push to context
        ContextProvider.Set(ContextProvider.Keys.ArrInstanceUrl, instance.ExternalOrInternalUrl);
        ContextProvider.Set(nameof(InstanceType), instance.ArrConfig.Type);
        ContextProvider.Set(ContextProvider.Keys.ArrInstanceId, instance.Id);
        ContextProvider.Set(ContextProvider.Keys.Version, instance.Version);

        IReadOnlyList<IDownloadService> downloadServices = await GetInitializedDownloadServicesAsync();

        ContentBlockerConfig config = ContextProvider.Get<ContentBlockerConfig>();

        bool targetResolved = false;

        await _arrArrQueueIterator.Iterate(arrClient, instance, async items =>
        {
            var groups = items
                .GroupBy(x => x.DownloadId)
                .ToList();

            foreach (var group in groups)
            {
                QueueRecord record = group.First();

                if (!arrClient.IsRecordValid(record))
                {
                    continue;
                }

                if (target is not null &&
                    !string.Equals(record.DownloadId, target.DownloadId, StringComparison.InvariantCultureIgnoreCase))
                {
                    continue;
                }

                bool resolved = await TryProcessRecordAsync(group, instance, arrClient, downloadServices, ignoredDownloads, config);

                if (target is not null)
                {
                    targetResolved = resolved;
                }
            }
        }, contentId: target is { ContentId: > 0 } ? target.ContentId : null);

        return targetResolved;
    }

    /// <summary>
    /// Scans a single grouped download. Returns <c>true</c> when the download is resolved — found in a
    /// client and evaluated, a usenet record, or deliberately skipped — and <c>false</c> only when it is
    /// a torrent that was not found in any download client yet (the case a webhook scan should retry).
    /// </summary>
    private async Task<bool> TryProcessRecordAsync(
        IGrouping<string, QueueRecord> group,
        ArrInstance instance,
        IArrClient arrClient,
        IReadOnlyList<IDownloadService> downloadServices,
        List<string> ignoredDownloads,
        ContentBlockerConfig config)
    {
        QueueRecord record = group.First();

        if (ignoredDownloads.Contains(record.DownloadId, StringComparer.InvariantCultureIgnoreCase))
        {
            _logger.LogInformation("skip | {title} | ignored", record.Title);
            return true;
        }

        _logger.LogTrace("processing | {title} | {id}", record.Title, record.DownloadId);

        bool hasContentId = arrClient.HasContentId(record);

        if (!hasContentId)
        {
            if (!config.ProcessNoContentId)
            {
                _logger.LogInformation("skip | item is missing the content id | {title}", record.Title);
                return true;
            }

            _logger.LogDebug("item is missing the content id | {title}", record.Title);
        }

        string downloadRemovalKey = CacheKeys.DownloadMarkedForRemoval(record.DownloadId, instance.Url);

        if (_cache.TryGetValue(downloadRemovalKey, out bool _))
        {
            _logger.LogDebug("skip | already marked for removal | {title}", record.Title);
            return true;
        }

        // push record to context
        ContextProvider.Set(nameof(QueueRecord), record);

        bool isTorrent = record.Protocol.Contains("torrent", StringComparison.InvariantCultureIgnoreCase);

        if (!isTorrent)
        {
            // Usenet is acknowledged once it appears in the queue; nothing to scan, no retry.
            return true;
        }

        BlockFilesResult result = new();
        DownloadClientConfig? foundInClient = null;

        var torrentClients = downloadServices
            .Where(x => x.ClientConfig.Type is DownloadClientType.Torrent)
            .ToList();

        _logger.LogDebug("searching unwanted files for {title}", record.Title);

        if (torrentClients.Count > 0)
        {
            // Check each download client for the download item
            foreach (var downloadService in torrentClients)
            {
                try
                {
                    result = await downloadService
                        .BlockUnwantedFilesAsync(record.DownloadId, ignoredDownloads);

                    if (result.Found)
                    {
                        foundInClient = downloadService.ClientConfig;
                        break;
                    }
                }
                catch (Exception ex)
                {
                    _logger.LogError(ex, "Error checking download {dName} with download client {cName}",
                        record.Title, downloadService.ClientConfig.Name);
                }
            }

            if (!result.Found)
            {
                _logger.LogWarning("Download not found in any torrent client | {title}", record.Title);
            }
        }
        else
        {
            _logger.LogDebug("No torrent clients enabled");
        }

        if (!result.Found || !result.MetadataFound)
        {
            // Retry while the torrent is not yet in a client, or is present but its file list/metadata isn't ready.
            return false;
        }

        if (result.ShouldRemove)
        {
            bool removeFromClient = true;

            if (result.IsPrivate && !config.DeletePrivate)
            {
                removeFromClient = false;
            }

            await PublishQueueItemRemoveRequest(
                downloadRemovalKey,
                instance,
                record,
                group.Count() > 1,
                removeFromClient,
                result.DeleteReason,
                skipSearch: !hasContentId,
                downloadClient: foundInClient
            );
        }

        return true;
    }
}