MalwareMultiScan/MalwareMultiScan.Api/Services/ScanResultsService.cs

using System;
using System.IO;
using System.Net.Http;
using System.Threading;
using System.Threading.Tasks;
using MalwareMultiScan.Api.Data.Configuration;
using MalwareMultiScan.Api.Data.Models;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.Logging;
using MongoDB.Bson;
using MongoDB.Driver;
using MongoDB.Driver.GridFS;

namespace MalwareMultiScan.Api.Services
{
    public class ScanResultsService
    {
        private const string CollectionName = "ScanResults";
        
        private readonly IMongoCollection<ScanResult> _collection;
        private readonly GridFSBucket _bucket;
        
        private readonly ScanBackendService _scanBackendService;

        public ScanResultsService(IMongoDatabase db, ScanBackendService scanBackendService)
        {
            _scanBackendService = scanBackendService;
            
            _collection = db.GetCollection<ScanResult>(CollectionName);
            _bucket = new GridFSBucket(db);
        }

        public async Task<ScanResult> CreateScanResult()
        {
            var scanResult = new ScanResult();

            await _collection.InsertOneAsync(scanResult);
            
            return scanResult;
        }
        
        public async Task UpdateScanResultForBackend(
            string resultId, string backendId, bool completed, bool succeeded, string[] threats = null)
        {
            var filterScanResult = Builders<ScanResult>.Filter.Where(r => r.Id == resultId);
            
            var updateScanResult = Builders<ScanResult>.Update.Set(r => r.Results[backendId], new ScanResultEntry
            {
                Completed = completed,
                Succeeded = succeeded,
                Threats = threats
            });

            await _collection.UpdateOneAsync(filterScanResult, updateScanResult);
        }

        public async Task QueueFileScan(ScanResult result, string fileName, Stream fileStream, string callbackUrl)
        {
            foreach (var backend in _scanBackendService.List)
            {
                var queueResult = await _scanBackendService.QueueFileScan(
                    backend, fileName, fileStream, callbackUrl);
                
                await UpdateScanResultForBackend(result.Id, backend.Id, !queueResult, queueResult);
            }
        }
        
        public async Task QueueUrlScan(ScanResult result, string url, string callbackUrl)
        {
            foreach (var backend in _scanBackendService.List)
            {
                var queueResult = await _scanBackendService.QueueUrlScan(
                    backend, url, callbackUrl);
                
                await UpdateScanResultForBackend(result.Id, backend.Id, !queueResult, queueResult);
            }
        }

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

            return objectId.ToString();
        }

        public async Task<Stream> ObtainFile(string id)
        {
            if (!ObjectId.TryParse(id, out var objectId))
                return null;
            
            return await _bucket.OpenDownloadStreamAsync(objectId, new GridFSDownloadOptions
            {
                Seekable = true
            });
        }
    }
}