The next generation of the Teknik Services. Written in ASP.NET. https://www.teknik.io/
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

Program.cs 13KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315
  1. using CommandLine;
  2. using Microsoft.EntityFrameworkCore;
  3. using nClam;
  4. using System;
  5. using System.Collections.Generic;
  6. using System.IO;
  7. using System.Linq;
  8. using System.Reflection;
  9. using System.Text;
  10. using System.Threading;
  11. using System.Threading.Tasks;
  12. using Teknik.Areas.Paste.Models;
  13. using Teknik.Areas.Stats.Models;
  14. using Teknik.Areas.Upload.Models;
  15. using Teknik.Areas.Users.Models;
  16. using Teknik.Areas.Users.Utility;
  17. using Teknik.Configuration;
  18. using Teknik.Data;
  19. using Teknik.Utilities;
  20. using Teknik.Utilities.Cryptography;
  21. namespace Teknik.ServiceWorker
  22. {
  23. public class Program
  24. {
  25. private static string currentPath = Path.GetDirectoryName(Assembly.GetEntryAssembly().Location);
  26. private static string virusFile = Path.Combine(currentPath, "virusLogs.txt");
  27. private static string errorFile = Path.Combine(currentPath, "errorLogs.txt");
  28. private static string orphansFile = Path.Combine(currentPath, "orphanedFiles.txt");
  29. private static string configPath = currentPath;
  30. private const string TAKEDOWN_REPORTER = "Teknik Automated System";
  31. private static readonly object dbLock = new object();
  32. private static readonly object scanStatsLock = new object();
  33. public static event Action<string> OutputEvent;
  34. public static int Main(string[] args)
  35. {
  36. try
  37. {
  38. Parser.Default.ParseArguments<ArgumentOptions>(args).WithParsed(options =>
  39. {
  40. if (!string.IsNullOrEmpty(options.Config))
  41. configPath = options.Config;
  42. if (Directory.Exists(configPath))
  43. {
  44. Config config = Config.Load(configPath);
  45. Output(string.Format("[{0}] Started Server Maintenance Process.", DateTime.Now));
  46. var optionsBuilder = new DbContextOptionsBuilder<TeknikEntities>();
  47. optionsBuilder.UseSqlServer(config.DbConnection);
  48. using (TeknikEntities db = new TeknikEntities(optionsBuilder.Options))
  49. {
  50. // Scan all the uploads for viruses, and remove the bad ones
  51. if (options.ScanUploads && config.UploadConfig.VirusScanEnable)
  52. {
  53. ScanUploads(config, db);
  54. }
  55. // Runs the migration
  56. if (options.Migrate)
  57. {
  58. // Run the overall migration calls
  59. TeknikMigration.RunMigration(db, config);
  60. }
  61. if (options.Expire)
  62. {
  63. ProcessExpirations(config, db);
  64. }
  65. if (options.Clean)
  66. {
  67. CleanStorage(config, db);
  68. }
  69. }
  70. Output(string.Format("[{0}] Finished Server Maintenance Process.", DateTime.Now));
  71. }
  72. else
  73. {
  74. string msg = string.Format("[{0}] Config File does not exist.", DateTime.Now);
  75. File.AppendAllLines(errorFile, new List<string> { msg });
  76. Output(msg);
  77. }
  78. });
  79. }
  80. catch (Exception ex)
  81. {
  82. string msg = string.Format("[{0}] Exception: {1}", DateTime.Now, ex.GetFullMessage(true));
  83. File.AppendAllLines(errorFile, new List<string> { msg });
  84. Output(msg);
  85. }
  86. return -1;
  87. }
  88. public static void ScanUploads(Config config, TeknikEntities db)
  89. {
  90. Output(string.Format("[{0}] Started Virus Scan.", DateTime.Now));
  91. List<Upload> uploads = db.Uploads.ToList();
  92. int maxConcurrency = 100;
  93. int totalCount = uploads.Count();
  94. int totalScans = 0;
  95. int currentScan = 0;
  96. int totalViruses = 0;
  97. using (SemaphoreSlim concurrencySemaphore = new SemaphoreSlim(maxConcurrency))
  98. {
  99. List<Task> runningTasks = new List<Task>();
  100. foreach (Upload upload in uploads)
  101. {
  102. concurrencySemaphore.Wait();
  103. currentScan++;
  104. Task scanTask = Task.Factory.StartNew(async () =>
  105. {
  106. try
  107. {
  108. var virusDetected = await ScanUpload(config, db, upload, totalCount, currentScan);
  109. if (virusDetected)
  110. totalViruses++;
  111. totalScans++;
  112. }
  113. catch (Exception ex)
  114. {
  115. string errorMsg = string.Format("[{0}] Scan Error: {1}", DateTime.Now, ex.GetFullMessage(true, true));
  116. File.AppendAllLines(errorFile, new List<string> { errorMsg });
  117. }
  118. finally
  119. {
  120. concurrencySemaphore.Release();
  121. }
  122. });
  123. if (scanTask != null)
  124. {
  125. runningTasks.Add(scanTask);
  126. }
  127. }
  128. Task.WaitAll(runningTasks.ToArray());
  129. }
  130. bool running = true;
  131. //while (running)
  132. //{
  133. // running = runningTasks.Exists(s => s != null && !s.IsCompleted && !s.IsCanceled && !s.IsFaulted);
  134. //}
  135. Output(string.Format("Scanning Complete. {0} Scanned | {1} Viruses Found | {2} Total Files", totalScans, totalViruses, totalCount));
  136. }
  137. private static async Task<bool> ScanUpload(Config config, TeknikEntities db, Upload upload, int totalCount, int currentCount)
  138. {
  139. bool virusDetected = false;
  140. string subDir = upload.FileName[0].ToString();
  141. string filePath = Path.Combine(config.UploadConfig.UploadDirectory, subDir, upload.FileName);
  142. if (File.Exists(filePath))
  143. {
  144. // If the IV is set, and Key is set, then scan it
  145. if (!string.IsNullOrEmpty(upload.Key) && !string.IsNullOrEmpty(upload.IV))
  146. {
  147. byte[] keyBytes = Encoding.UTF8.GetBytes(upload.Key);
  148. byte[] ivBytes = Encoding.UTF8.GetBytes(upload.IV);
  149. long maxUploadSize = config.UploadConfig.MaxUploadSize;
  150. if (upload.User != null)
  151. {
  152. maxUploadSize = config.UploadConfig.MaxUploadSizeBasic;
  153. IdentityUserInfo userInfo = await IdentityHelper.GetIdentityUserInfo(config, upload.User.Username);
  154. if (userInfo.AccountType == AccountType.Premium)
  155. {
  156. maxUploadSize = config.UploadConfig.MaxUploadSizePremium;
  157. }
  158. }
  159. using (FileStream fs = new FileStream(filePath, FileMode.Open, FileAccess.Read))
  160. using (AesCounterStream aesStream = new AesCounterStream(fs, false, keyBytes, ivBytes))
  161. {
  162. ClamClient clam = new ClamClient(config.UploadConfig.ClamServer, config.UploadConfig.ClamPort);
  163. clam.MaxStreamSize = maxUploadSize;
  164. ClamScanResult scanResult = await clam.SendAndScanFileAsync(fs);
  165. switch (scanResult.Result)
  166. {
  167. case ClamScanResults.Clean:
  168. string cleanMsg = string.Format("[{0}] Clean Scan: {1}/{2} Scanned | {3} - {4}", DateTime.Now, currentCount, totalCount, upload.Url, upload.FileName);
  169. Output(cleanMsg);
  170. break;
  171. case ClamScanResults.VirusDetected:
  172. string msg = string.Format("[{0}] Virus Detected: {1} - {2} - {3}", DateTime.Now, upload.Url, upload.FileName, scanResult.InfectedFiles.First().VirusName);
  173. Output(msg);
  174. lock (scanStatsLock)
  175. {
  176. virusDetected = true;
  177. File.AppendAllLines(virusFile, new List<string> { msg });
  178. }
  179. lock (dbLock)
  180. {
  181. string urlName = upload.Url;
  182. // Delete from the DB
  183. db.Uploads.Remove(upload);
  184. // Delete the File
  185. if (File.Exists(filePath))
  186. {
  187. File.Delete(filePath);
  188. }
  189. // Add to transparency report if any were found
  190. Takedown report = new Takedown();
  191. report.Requester = TAKEDOWN_REPORTER;
  192. report.RequesterContact = config.SupportEmail;
  193. report.DateRequested = DateTime.Now;
  194. report.Reason = "Malware Found";
  195. report.ActionTaken = string.Format("Upload removed: {0}", urlName);
  196. report.DateActionTaken = DateTime.Now;
  197. db.Takedowns.Add(report);
  198. // Save Changes
  199. db.SaveChanges();
  200. }
  201. break;
  202. case ClamScanResults.Error:
  203. string errorMsg = string.Format("[{0}] Scan Error: {1}", DateTime.Now, scanResult.RawResult);
  204. File.AppendAllLines(errorFile, new List<string> { errorMsg });
  205. Output(errorMsg);
  206. break;
  207. case ClamScanResults.Unknown:
  208. string unkMsg = string.Format("[{0}] Unknown Scan Result: {1}", DateTime.Now, scanResult.RawResult);
  209. File.AppendAllLines(errorFile, new List<string> { unkMsg });
  210. Output(unkMsg);
  211. break;
  212. }
  213. }
  214. }
  215. }
  216. return virusDetected;
  217. }
  218. public static void ProcessExpirations(Config config, TeknikEntities db)
  219. {
  220. Output(string.Format("[{0}] Starting processing expirations.", DateTime.Now));
  221. var curDate = DateTime.Now;
  222. // Process uploads
  223. List<Upload> uploads = db.Uploads.Where(u => u.ExpireDate != null && u.ExpireDate < curDate).ToList();
  224. foreach (Upload upload in uploads)
  225. {
  226. string subDir = upload.FileName[0].ToString();
  227. string filePath = Path.Combine(config.UploadConfig.UploadDirectory, subDir, upload.FileName);
  228. // Delete the File
  229. if (File.Exists(filePath))
  230. {
  231. File.Delete(filePath);
  232. }
  233. }
  234. db.RemoveRange(uploads);
  235. db.SaveChanges();
  236. // Process Pastes
  237. List<Paste> pastes = db.Pastes.Where(p => p.ExpireDate != null && p.ExpireDate < curDate).ToList();
  238. foreach (Paste paste in pastes)
  239. {
  240. string subDir = paste.FileName[0].ToString();
  241. string filePath = Path.Combine(config.PasteConfig.PasteDirectory, subDir, paste.FileName);
  242. // Delete the File
  243. if (File.Exists(filePath))
  244. {
  245. File.Delete(filePath);
  246. }
  247. }
  248. db.RemoveRange(pastes);
  249. db.SaveChanges();
  250. }
  251. public static void CleanStorage(Config config, TeknikEntities db)
  252. {
  253. var curDate = DateTime.Now;
  254. // Process upload data
  255. Output(string.Format("[{0}] Starting processing upload storage cleaning.", DateTime.Now));
  256. List<string> uploads = db.Uploads.Select(u => Path.Combine(config.UploadConfig.UploadDirectory, u.FileName[0].ToString(), u.FileName)).Select(u => u.ToLower()).ToList();
  257. List<string> files = Directory.GetFiles(config.UploadConfig.UploadDirectory, "*.*", SearchOption.AllDirectories).Select(f => f.ToLower()).ToList();
  258. var orphans = files.Except(uploads);
  259. File.AppendAllLines(orphansFile, orphans);
  260. foreach (var orphan in orphans)
  261. {
  262. File.Delete(orphan);
  263. }
  264. }
  265. public static void Output(string message)
  266. {
  267. Console.WriteLine(message);
  268. if (OutputEvent != null)
  269. {
  270. OutputEvent(message);
  271. }
  272. }
  273. }
  274. }