From 4207a2811fe565f925555749c492b31796377a40 Mon Sep 17 00:00:00 2001
From: Uncled1023 <admin@teknik.io>
Date: Tue, 31 Dec 2019 00:34:46 -0800
Subject: [PATCH] Added storage cleaning routine and fixed upload virus
 scanning.

---
 ServiceWorker/ArgumentOptions.cs              |  3 +
 ServiceWorker/Program.cs                      | 77 +++++++++++++++----
 .../Teknik Service Worker.pubxml              |  3 +-
 3 files changed, 67 insertions(+), 16 deletions(-)

diff --git a/ServiceWorker/ArgumentOptions.cs b/ServiceWorker/ArgumentOptions.cs
index 9e22277..cf4bb83 100644
--- a/ServiceWorker/ArgumentOptions.cs
+++ b/ServiceWorker/ArgumentOptions.cs
@@ -22,6 +22,9 @@ namespace Teknik.ServiceWorker
         [Option('e', "expire", Default = false, Required = false, HelpText = "Process Expirations")]
         public bool Expire { get; set; }
 
+        [Option('C', "clean", Default = false, Required = false, HelpText = "Clean Storage")]
+        public bool Clean { get; set; }
+
         // Omitting long name, default --verbose
         [Option(HelpText = "Prints all messages to standard output.")]
         public bool Verbose { get; set; }
diff --git a/ServiceWorker/Program.cs b/ServiceWorker/Program.cs
index dd259f6..6c94039 100644
--- a/ServiceWorker/Program.cs
+++ b/ServiceWorker/Program.cs
@@ -7,6 +7,7 @@ using System.IO;
 using System.Linq;
 using System.Reflection;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Teknik.Areas.Paste.Models;
 using Teknik.Areas.Stats.Models;
@@ -25,6 +26,7 @@ namespace Teknik.ServiceWorker
         private static string currentPath = Path.GetDirectoryName(Assembly.GetEntryAssembly().Location);
         private static string virusFile = Path.Combine(currentPath, "virusLogs.txt");
         private static string errorFile = Path.Combine(currentPath, "errorLogs.txt");
+        private static string orphansFile = Path.Combine(currentPath, "orphanedFiles.txt");
         private static string configPath = currentPath;
 
         private const string TAKEDOWN_REPORTER = "Teknik Automated System";
@@ -70,6 +72,11 @@ namespace Teknik.ServiceWorker
                             {
                                 ProcessExpirations(config, db);
                             }
+
+                            if (options.Clean)
+                            {
+                                CleanStorage(config, db);
+                            }
                         }
 
                         Output(string.Format("[{0}] Finished Server Maintenance Process.", DateTime.Now));
@@ -96,30 +103,53 @@ namespace Teknik.ServiceWorker
             Output(string.Format("[{0}] Started Virus Scan.", DateTime.Now));
             List<Upload> uploads = db.Uploads.ToList();
 
+            int maxConcurrency = 100;
             int totalCount = uploads.Count();
             int totalScans = 0;
+            int currentScan = 0;
             int totalViruses = 0;
-            List<Task> runningTasks = new List<Task>();
-            foreach (Upload upload in uploads)
+
+            using (SemaphoreSlim concurrencySemaphore = new SemaphoreSlim(maxConcurrency))
             {
-                int currentScan = totalScans++;
-                Task scanTask = Task.Factory.StartNew(async () =>
+                List<Task> runningTasks = new List<Task>();
+                foreach (Upload upload in uploads)
                 {
-                    var virusDetected = await ScanUpload(config, db, upload, totalCount, currentScan);
-                    if (virusDetected)
-                        totalViruses++;
-                });
-                if (scanTask != null)
-                {
-                    runningTasks.Add(scanTask);
+                    concurrencySemaphore.Wait();
+
+                    currentScan++;
+
+                    Task scanTask = Task.Factory.StartNew(async () =>
+                    {
+                        try
+                        {
+                            var virusDetected = await ScanUpload(config, db, upload, totalCount, currentScan);
+                            if (virusDetected)
+                                totalViruses++;
+                            totalScans++;
+                        }
+                        catch (Exception ex)
+                        {
+                            string errorMsg = string.Format("[{0}] Scan Error: {1}", DateTime.Now, ex.GetFullMessage(true, true));
+                            File.AppendAllLines(errorFile, new List<string> { errorMsg });
+                        }
+                        finally
+                        {
+                            concurrencySemaphore.Release();
+                        }
+                    });
+                    if (scanTask != null)
+                    {
+                        runningTasks.Add(scanTask);
+                    }
                 }
+                Task.WaitAll(runningTasks.ToArray());
             }
             bool running = true;
 
-            while (running)
-            {
-                running = runningTasks.Exists(s => s != null && !s.IsCompleted && !s.IsCanceled && !s.IsFaulted);
-            }
+            //while (running)
+            //{
+            //    running = runningTasks.Exists(s => s != null && !s.IsCompleted && !s.IsCanceled && !s.IsFaulted);
+            //}
 
             Output(string.Format("Scanning Complete.  {0} Scanned | {1} Viruses Found | {2} Total Files", totalScans, totalViruses, totalCount));
         }
@@ -256,6 +286,23 @@ namespace Teknik.ServiceWorker
             db.SaveChanges();
         }
 
+        public static void CleanStorage(Config config, TeknikEntities db)
+        {
+            var curDate = DateTime.Now;
+
+            // Process upload data
+            Output(string.Format("[{0}] Starting processing upload storage cleaning.", DateTime.Now));
+
+            List<string> uploads = db.Uploads.Select(u => Path.Combine(config.UploadConfig.UploadDirectory, u.FileName[0].ToString(), u.FileName)).Select(u => u.ToLower()).ToList();
+            List<string> files = Directory.GetFiles(config.UploadConfig.UploadDirectory, "*.*", SearchOption.AllDirectories).Select(f => f.ToLower()).ToList();
+            var orphans = files.Except(uploads);
+            File.AppendAllLines(orphansFile, orphans);
+            foreach (var orphan in orphans)
+            {
+                File.Delete(orphan);
+            }
+        }
+
         public static void Output(string message)
         {
             Console.WriteLine(message);
diff --git a/ServiceWorker/Properties/PublishProfiles/Teknik Service Worker.pubxml b/ServiceWorker/Properties/PublishProfiles/Teknik Service Worker.pubxml
index 18047f3..25a249e 100644
--- a/ServiceWorker/Properties/PublishProfiles/Teknik Service Worker.pubxml	
+++ b/ServiceWorker/Properties/PublishProfiles/Teknik Service Worker.pubxml	
@@ -7,7 +7,8 @@ https://go.microsoft.com/fwlink/?LinkID=208121.
     <PublishProtocol>FileSystem</PublishProtocol>
     <Configuration>Release</Configuration>
     <Platform>Any CPU</Platform>
-    <TargetFramework>netcoreapp2.1</TargetFramework>
+    <TargetFramework>netcoreapp2.2</TargetFramework>
     <PublishDir>Output\</PublishDir>
+    <SelfContained>false</SelfContained>
   </PropertyGroup>
 </Project>
\ No newline at end of file