MalwareMultiScan/MalwareMultiScan.Api/Services/ScanResultService.cs

using System.IO;
using System.Linq;
using System.Threading.Tasks;
using MalwareMultiScan.Api.Data.Models;
using MongoDB.Bson;
using MongoDB.Driver;
using MongoDB.Driver.GridFS;

namespace MalwareMultiScan.Api.Services
{
    /// <summary>
    /// Scan results service.
    /// </summary>
    public class ScanResultService
    {
        private const string CollectionName = "ScanResults";

        private readonly GridFSBucket _bucket;
        private readonly IMongoCollection<ScanResult> _collection;
        private readonly ScanBackendService _scanBackendService;

        /// <summary>
        /// Create scan result service.
        /// </summary>
        /// <param name="db">Mongo database.</param>
        /// <param name="scanBackendService">Scan backend service.</param>
        public ScanResultService(IMongoDatabase db, ScanBackendService scanBackendService)
        {
            _scanBackendService = scanBackendService;

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

        /// <summary>
        /// Create scan result.
        /// </summary>
        /// <returns>Scan result.</returns>
        public async Task<ScanResult> CreateScanResult()
        {
            var scanResult = new ScanResult
            {
                Results = _scanBackendService.List
                    .Where(b => b.Enabled)
                    .ToDictionary(k => k.Id, v => new ScanResultEntry())
            };

            await _collection.InsertOneAsync(scanResult);

            return scanResult;
        }

        /// <summary>
        /// Get scan result.
        /// </summary>
        /// <param name="id">Scan result id.</param>
        /// <returns>Scan result.</returns>
        public async Task<ScanResult> GetScanResult(string id)
        {
            var result = await _collection.FindAsync(
                Builders<ScanResult>.Filter.Where(r => r.Id == id));

            return await result.FirstOrDefaultAsync();
        }

        /// <summary>
        /// Update scan status for the backend.
        /// </summary>
        /// <param name="resultId">Result id.</param>
        /// <param name="backendId">Backend id.</param>
        /// <param name="duration">Duration of scan.</param>
        /// <param name="completed">If the scan has been completed.</param>
        /// <param name="succeeded">If the scan has been succeeded.</param>
        /// <param name="threats">List of found threats.</param>
        public async Task UpdateScanResultForBackend(string resultId, string backendId, long duration,
            bool completed = false, bool succeeded = false, string[] threats = null)
        {
            await _collection.UpdateOneAsync(
                Builders<ScanResult>.Filter.Where(r => r.Id == resultId),
                Builders<ScanResult>.Update.Set(r => r.Results[backendId], new ScanResultEntry
                {
                    Completed = completed,
                    Succeeded = succeeded,
                    Duration = duration,
                    Threats = threats ?? new string[] { }
                }));
        }

        /// <summary>
        /// Queue URL scan.
        /// </summary>
        /// <param name="result">Scan result instance.</param>
        /// <param name="fileUrl">File URL.</param>
        public async Task QueueUrlScan(ScanResult result, string fileUrl)
        {
            foreach (var backend in _scanBackendService.List.Where(b => b.Enabled))
                await _scanBackendService.QueueUrlScan(result, backend, fileUrl);
        }

        /// <summary>
        /// Store file.
        /// </summary>
        /// <param name="fileName">File name.</param>
        /// <param name="fileStream">File stream.</param>
        /// <returns>Stored file id.</returns>
        public async Task<string> StoreFile(string fileName, Stream fileStream)
        {
            var objectId = await _bucket.UploadFromStreamAsync(
                fileName, fileStream);

            return objectId.ToString();
        }

        /// <summary>
        /// Obtain stored file stream.
        /// </summary>
        /// <param name="id">File id.</param>
        /// <returns>File seekable stream.</returns>
        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
            });
        }
    }
}