MalwareMultiScan/MalwareMultiScan.Api/Services/ScanResultService.cs

using System;
using System.IO;
using System.Linq;
using System.Threading.Tasks;
using Consul;
using Hangfire;
using Hangfire.States;
using MalwareMultiScan.Api.Data;
using MalwareMultiScan.Api.Services.Interfaces;
using MalwareMultiScan.Shared.Enums;
using MalwareMultiScan.Shared.Message;
using MalwareMultiScan.Shared.Services.Interfaces;
using MongoDB.Bson;
using MongoDB.Driver;
using MongoDB.Driver.GridFS;

namespace MalwareMultiScan.Api.Services
{
    /// <inheritdoc />
    public class ScanResultService : IScanResultService
    {
        private const string CollectionName = "ScanResults";
        private readonly IBackgroundJobClient _backgroundJobClient;

        private readonly IGridFSBucket _bucket;
        private readonly IMongoCollection<ScanResult> _collection;
        private readonly IConsulClient _consulClient;

        /// <summary>
        /// Initialize scan result service.
        /// </summary>
        /// <param name="db">Mongo database.</param>
        /// <param name="bucket">GridFS bucket.</param>
        /// <param name="consulClient">Consul client.</param>
        /// <param name="backgroundJobClient">Background job client.</param>
        public ScanResultService(
            IMongoDatabase db,
            IGridFSBucket bucket,
            IConsulClient consulClient,
            IBackgroundJobClient backgroundJobClient)
        {
            _bucket = bucket;
            _consulClient = consulClient;
            _backgroundJobClient = backgroundJobClient;

            _collection = db.GetCollection<ScanResult>(CollectionName);
        }

        /// <inheritdoc />
        public async Task<ScanResult> CreateScanResult(Uri callbackUrl)
        {
            var scanResult = new ScanResult
            {
                CallbackUrl = callbackUrl
            };

            await _collection.InsertOneAsync(scanResult);

            return scanResult;
        }

        /// <inheritdoc />
        public async Task<ScanResult> GetScanResult(string id)
        {
            var result = await _collection.FindAsync(
                Builders<ScanResult>.Filter.Where(r => r.Id == id));

            return await result.FirstOrDefaultAsync();
        }

        /// <inheritdoc />
        public async Task UpdateScanResultForBackend(string resultId, string backendId,
            ScanResultMessage result = null)
        {
            result ??= new ScanResultMessage
            {
                Status = ScanResultStatus.Queued
            };

            await _collection.UpdateOneAsync(
                Builders<ScanResult>.Filter.Where(r => r.Id == resultId),
                Builders<ScanResult>.Update.Set(r => r.Results[backendId], result));
        }

        /// <inheritdoc />
        public async Task QueueUrlScan(ScanResult result, string fileUrl)
        {
            var message = new ScanQueueMessage
            {
                Id = result.Id,
                Uri = new Uri(fileUrl)
            };

            var scanners = await _consulClient.Health.Service("scanner", null, true);

            var backends = scanners.Response
                .Select(s => s.Service.Meta.TryGetValue("BackendId", out var backendId) ? backendId : null)
                .Where(q => q != null)
                .Distinct()
                .ToArray();

            foreach (var backend in backends)
            {
                await UpdateScanResultForBackend(result.Id, backend);

                _backgroundJobClient.Create<IScanBackgroundJob>(
                    j => j.Process(message), new EnqueuedState(backend));
            }
        }

        /// <inheritdoc />
        public async Task<string> StoreFile(string fileName, Stream fileStream)
        {
            var objectId = await _bucket.UploadFromStreamAsync(
                fileName, fileStream);

            return objectId.ToString();
        }

        /// <inheritdoc />
        public async Task<Stream> ObtainFile(string id)
        {
            if (!ObjectId.TryParse(id, out var objectId))
                return null;

            try
            {
                return await _bucket.OpenDownloadStreamAsync(objectId, new GridFSDownloadOptions
                {
                    Seekable = true
                });
            }
            catch (GridFSFileNotFoundException)
            {
                return null;
            }
        }
    }
}