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 11KB

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