From 45c0989251f45d5cfc0c6f9c5d4a430f2a10a675 Mon Sep 17 00:00:00 2001 From: p0rtL6 Date: Sun, 6 Apr 2025 19:08:44 -0700 Subject: [PATCH 1/3] Add pause/resume functionality --- .../Concurrency/BlockingTaskScheduler.cs | 176 +++++++++++++++++ SnaffCore/Config/Options.cs | 4 + SnaffCore/SnaffCon.cs | 186 +++++++++++------- SnaffCore/TreeWalk/TreeWalker.cs | 54 ++--- Snaffler/Config.cs | 14 ++ 5 files changed, 335 insertions(+), 99 deletions(-) diff --git a/SnaffCore/Concurrency/BlockingTaskScheduler.cs b/SnaffCore/Concurrency/BlockingTaskScheduler.cs index f748bd3a..876948f0 100644 --- a/SnaffCore/Concurrency/BlockingTaskScheduler.cs +++ b/SnaffCore/Concurrency/BlockingTaskScheduler.cs @@ -1,6 +1,9 @@ using System; using System.Collections.Generic; +using System.IO; +using System.Linq; using System.Numerics; +using System.Text; using System.Threading; using System.Threading.Tasks; @@ -34,6 +37,8 @@ public bool Done() Scheduler.RecalculateCounters(); TaskCounters taskCounters = Scheduler.GetTaskCounters(); + Console.WriteLine($"Checking if done - queued: {taskCounters.CurrentTasksQueued}, done: {taskCounters.CurrentTasksRunning}"); + if ((taskCounters.CurrentTasksQueued + taskCounters.CurrentTasksRunning == 0)) { return true; @@ -68,6 +73,177 @@ public void New(Action action) } } + public enum TaskType + { + None = 0, + Share = 1, + Tree = 2, + File = 3 + } + + public struct TaskFileEntry + { + public TaskFileEntryStatus status; + public string guid; + public TaskType type; + public string input; + + public TaskFileEntry(string entryLine) + { + string[] lineParts = entryLine.Split('|'); + + status = (TaskFileEntryStatus)Enum.Parse(typeof(TaskFileEntryStatus), lineParts[0]); + guid = lineParts[1]; + + if (status == TaskFileEntryStatus.Pending) + { + type = (TaskType)Enum.Parse(typeof(TaskType), lineParts[2]); + input = lineParts[3]; + } + else + { + type = TaskType.None; + input = null; + } + } + } + + public enum TaskFileEntryStatus + { + Pending = 0, + Completed = 1, + } + + public class ResumingTaskScheduler : BlockingStaticTaskScheduler + { + private static readonly object WriteLock = new object(); + private static StreamWriter fileWriter; + + private static string[] AlreadyHandledTasks; + + internal BlockingMq Mq { get; } + + public ResumingTaskScheduler(int threads, int maxBacklog) : base(threads, maxBacklog) + { + this.Mq = BlockingMq.GetMq(); + } + + public static void SetAlreadyHandledTasks(TaskFileEntry[] taskFileEntries) + { + AlreadyHandledTasks = taskFileEntries.Select(entry => entry.input).ToArray(); + } + + public static bool IsTaskAlreadyHandled(string input) + { + if (AlreadyHandledTasks == null) return false; + return AlreadyHandledTasks.Contains(input); + } + + public void New(TaskType taskType, Action action, string input) + { + if (IsTaskAlreadyHandled(input)) return; + + string taskGuid = SaveTask(taskType, input); + + New(() => + { + try + { + action(input); + } + catch (Exception e) + { + Mq.Error("Exception in " + taskType.ToString() + " task for host " + input); + Mq.Error(e.ToString()); + } + + CompleteTask(taskGuid); + }); + } + + public static void SetTaskFile(string path) + { + fileWriter = new StreamWriter(path); + } + + public static void CloseTaskFile() + { + if (fileWriter != null) { + lock (WriteLock) + { + fileWriter.Close(); + fileWriter = null; + } + } + } + + internal string SaveTask(TaskType taskType, string input) + { + // task file is not set, we are not saving tasks + if (fileWriter == null) return null; + + string taskGuid = Guid.NewGuid().ToString(); + + StringBuilder taskEntryBuilder = new StringBuilder(); + + taskEntryBuilder.Append("Pending|"); + taskEntryBuilder.Append(taskGuid); + taskEntryBuilder.Append("|"); + taskEntryBuilder.Append(taskType.ToString()); + taskEntryBuilder.Append("|"); + taskEntryBuilder.Append(input); + + lock (WriteLock) + { + fileWriter.WriteLine(taskEntryBuilder.ToString()); + fileWriter.Flush(); + } + + return taskGuid; + } + + internal void CompleteTask(string taskGuid) + { + if (fileWriter == null) return; + + lock (WriteLock) + { + fileWriter.WriteLine("Completed|" + taskGuid); + fileWriter.Flush(); + } + } + } + + public class ShareTaskScheduler : ResumingTaskScheduler + { + public ShareTaskScheduler(int threads, int maxBacklog) : base(threads, maxBacklog) { } + + public void New(Action action, string file) + { + New(TaskType.Share, action, file); + } + } + + public class TreeTaskScheduler : ResumingTaskScheduler + { + public TreeTaskScheduler(int threads, int maxBacklog) : base(threads, maxBacklog) { } + + public void New(Action action, string file) + { + New(TaskType.Tree, action, file); + } + } + + public class FileTaskScheduler : ResumingTaskScheduler + { + public FileTaskScheduler(int threads, int maxBacklog) : base(threads, maxBacklog) { } + + public void New(Action action, string file) + { + New(TaskType.File, action, file); + } + } + public class TaskCounters { public BigInteger TotalTasksQueued { get; set; } diff --git a/SnaffCore/Config/Options.cs b/SnaffCore/Config/Options.cs index 78c66282..4c6c6eec 100644 --- a/SnaffCore/Config/Options.cs +++ b/SnaffCore/Config/Options.cs @@ -14,6 +14,10 @@ public partial class Options { public static Options MyOptions { get; set; } + // Pause and resume functionality + public string TaskFile { get; set; } + public string ResumeFrom { get; set; } + // Manual Targeting Options public List PathTargets { get; set; } = new List(); public string[] ComputerTargets { get; set; } diff --git a/SnaffCore/SnaffCon.cs b/SnaffCore/SnaffCon.cs index 6caf1ab4..092c48f3 100644 --- a/SnaffCore/SnaffCon.cs +++ b/SnaffCore/SnaffCon.cs @@ -16,6 +16,8 @@ using static SnaffCore.Config.Options; using Timer = System.Timers.Timer; using System.Net; +using System.IO; +using Nett; namespace SnaffCore { @@ -25,10 +27,10 @@ public class SnaffCon private BlockingMq Mq { get; set; } - private static BlockingStaticTaskScheduler ShareTaskScheduler; - private static BlockingStaticTaskScheduler TreeTaskScheduler; - private static BlockingStaticTaskScheduler FileTaskScheduler; - + private static ShareTaskScheduler ShareTaskScheduler; + private static TreeTaskScheduler TreeTaskScheduler; + private static FileTaskScheduler FileTaskScheduler; + private static ShareFinder ShareFinder; private static TreeWalker TreeWalker; private static FileScanner FileScanner; @@ -46,10 +48,10 @@ public SnaffCon(Options options) int treeThreads = MyOptions.TreeThreads; int fileThreads = MyOptions.FileThreads; - ShareTaskScheduler = new BlockingStaticTaskScheduler(shareThreads, MyOptions.MaxShareQueue); - TreeTaskScheduler = new BlockingStaticTaskScheduler(treeThreads, MyOptions.MaxTreeQueue); - FileTaskScheduler = new BlockingStaticTaskScheduler(fileThreads, MyOptions.MaxFileQueue); - + ShareTaskScheduler = new ShareTaskScheduler(shareThreads, MyOptions.MaxShareQueue); + TreeTaskScheduler = new TreeTaskScheduler(treeThreads, MyOptions.MaxTreeQueue); + FileTaskScheduler = new FileTaskScheduler(fileThreads, MyOptions.MaxFileQueue); + FileScanner = new FileScanner(); TreeWalker = new TreeWalker(); ShareFinder = new ShareFinder(); @@ -67,15 +69,15 @@ public static FileScanner GetFileScanner() { return FileScanner; } - public static BlockingStaticTaskScheduler GetShareTaskScheduler() + public static ShareTaskScheduler GetShareTaskScheduler() { return ShareTaskScheduler; } - public static BlockingStaticTaskScheduler GetTreeTaskScheduler() + public static TreeTaskScheduler GetTreeTaskScheduler() { return TreeTaskScheduler; } - public static BlockingStaticTaskScheduler GetFileTaskScheduler() + public static FileTaskScheduler GetFileTaskScheduler() { return FileTaskScheduler; } @@ -91,72 +93,107 @@ public void Execute() statusUpdateTimer.Elapsed += TimedStatusUpdate; statusUpdateTimer.Start(); - - // If we want to hunt for user IDs, we need data from the running user's domain. - // Future - walk trusts - if ( MyOptions.DomainUserRules) + if (MyOptions.TaskFile != null && MyOptions.TaskFile.Length > 0) { - DomainUserDiscovery(); + ResumingTaskScheduler.SetTaskFile(MyOptions.TaskFile); } - // Explicit folder setting overrides DFS - if (MyOptions.PathTargets.Count != 0 && (MyOptions.DfsShareDiscovery || MyOptions.DfsOnly)) + if (MyOptions.ResumeFrom != null && MyOptions.ResumeFrom.Length > 0) { - DomainDfsDiscovery(); - } + // parse task file into structured data + TaskFileEntry[] taskFileEntries = File.ReadAllLines(MyOptions.ResumeFrom).Select(line => new TaskFileEntry(line)).ToArray(); + + TaskFileEntry[] completedTasks = taskFileEntries.Where(entry => entry.status == TaskFileEntryStatus.Completed).ToArray(); + TaskFileEntry[] pendingTasks = taskFileEntries.Where(entry => (entry.status == TaskFileEntryStatus.Pending && !completedTasks.Any(completedEntry => completedEntry.guid == entry.guid))).ToArray(); - if (MyOptions.PathTargets.Count == 0 && MyOptions.ComputerTargets == null) + foreach (TaskFileEntry pendingFileTask in pendingTasks) + { + switch (pendingFileTask.type) + { + case TaskType.Share: + ShareFinder shareFinder = new ShareFinder(); + ShareTaskScheduler.New(shareFinder.GetComputerShares, pendingFileTask.input); + break; + case TaskType.Tree: + TreeTaskScheduler.New(TreeWalker.WalkTree, pendingFileTask.input); + break; + case TaskType.File: + FileTaskScheduler.New(FileScanner.ScanFile, pendingFileTask.input); + break; + } + } + } + else { - if (MyOptions.DfsSharesDict.Count == 0) + // If we want to hunt for user IDs, we need data from the running user's domain. + // Future - walk trusts + if (MyOptions.DomainUserRules) { - Mq.Info("Invoking DFS Discovery because no ComputerTargets or PathTargets were specified"); - DomainDfsDiscovery(); + DomainUserDiscovery(); } - if (!MyOptions.DfsOnly) + // Explicit folder setting overrides DFS + if (MyOptions.PathTargets.Count != 0 && (MyOptions.DfsShareDiscovery || MyOptions.DfsOnly)) { - Mq.Info("Invoking full domain computer discovery."); - DomainTargetDiscovery(); + DomainDfsDiscovery(); } - else + + if (MyOptions.PathTargets.Count == 0 && MyOptions.ComputerTargets == null) { - Mq.Info("Skipping domain computer discovery."); - foreach (string share in MyOptions.DfsSharesDict.Keys) + if (MyOptions.DfsSharesDict.Count == 0) + { + Mq.Info("Invoking DFS Discovery because no ComputerTargets or PathTargets were specified"); + DomainDfsDiscovery(); + } + + if (!MyOptions.DfsOnly) { - if (!MyOptions.PathTargets.Contains(share)) + Mq.Info("Invoking full domain computer discovery."); + DomainTargetDiscovery(); + } + else + { + Mq.Info("Skipping domain computer discovery."); + foreach (string share in MyOptions.DfsSharesDict.Keys) { - MyOptions.PathTargets.Add(share); + if (!MyOptions.PathTargets.Contains(share)) + { + MyOptions.PathTargets.Add(share); + } } + Mq.Info("Starting TreeWalker tasks on DFS shares."); + FileDiscovery(MyOptions.PathTargets.ToArray()); } - Mq.Info("Starting TreeWalker tasks on DFS shares."); + } + // otherwise we should have a set of path targets... + else if (MyOptions.PathTargets.Count != 0) + { FileDiscovery(MyOptions.PathTargets.ToArray()); } - } - // otherwise we should have a set of path targets... - else if (MyOptions.PathTargets.Count != 0) - { - FileDiscovery(MyOptions.PathTargets.ToArray()); - } - // or we've been told what computers to hit... - else if (MyOptions.ComputerTargets != null) - { - ShareDiscovery(MyOptions.ComputerTargets); - } + // or we've been told what computers to hit... + else if (MyOptions.ComputerTargets != null) + { + ShareDiscovery(MyOptions.ComputerTargets); + } - // but if that hasn't been done, something has gone wrong. - else - { - Mq.Error("OctoParrot says: AWK! I SHOULDN'T BE!"); + // but if that hasn't been done, something has gone wrong. + else + { + Mq.Error("OctoParrot says: AWK! I SHOULDN'T BE!"); + } } waitHandle.WaitOne(); + ResumingTaskScheduler.CloseTaskFile(); StatusUpdate(); + DateTime finished = DateTime.Now; TimeSpan runSpan = finished.Subtract(StartTime); Mq.Info("Finished at " + finished.ToLocalTime()); Mq.Info("Snafflin' took " + runSpan); Mq.Finish(); + } private void DomainDfsDiscovery() @@ -297,19 +334,21 @@ private void ShareDiscovery(string[] computerTargets) } // ShareFinder Task Creation - this kicks off the rest of the flow Mq.Trace("Creating a ShareFinder task for " + computer); - ShareTaskScheduler.New(() => - { - try - { - ShareFinder shareFinder = new ShareFinder(); - shareFinder.GetComputerShares(computer); - } - catch (Exception e) - { - Mq.Error("Exception in ShareFinder task for host " + computer); - Mq.Error(e.ToString()); - } - }); + ShareFinder shareFinder = new ShareFinder(); + ShareTaskScheduler.New(shareFinder.GetComputerShares, computer); + //ShareTaskScheduler.New(() => + //{ + // try + // { + // ShareFinder shareFinder = new ShareFinder(); + // shareFinder.GetComputerShares(computer); + // } + // catch (Exception e) + // { + // Mq.Error("Exception in ShareFinder task for host " + computer); + // Mq.Error(e.ToString()); + // } + //}); } Mq.Info("Created all sharefinder tasks."); } @@ -371,18 +410,19 @@ private void FileDiscovery(string[] pathTargets) { // TreeWalker Task Creation - this kicks off the rest of the flow Mq.Info("Creating a TreeWalker task for " + pathTarget); - TreeTaskScheduler.New(() => - { - try - { - TreeWalker.WalkTree(pathTarget); - } - catch (Exception e) - { - Mq.Error("Exception in TreeWalker task for path " + pathTarget); - Mq.Error(e.ToString()); - } - }); + TreeTaskScheduler.New(TreeWalker.WalkTree, pathTarget); + //TreeTaskScheduler.New(() => + //{ + // try + // { + // TreeWalker.WalkTree(pathTarget); + // } + // catch (Exception e) + // { + // Mq.Error("Exception in TreeWalker task for path " + pathTarget); + // Mq.Error(e.ToString()); + // } + //}); } Mq.Info("Created all TreeWalker tasks."); diff --git a/SnaffCore/TreeWalk/TreeWalker.cs b/SnaffCore/TreeWalk/TreeWalker.cs index c3cf3fc3..825c90f3 100644 --- a/SnaffCore/TreeWalk/TreeWalker.cs +++ b/SnaffCore/TreeWalk/TreeWalker.cs @@ -10,8 +10,8 @@ namespace SnaffCore.TreeWalk public class TreeWalker { private BlockingMq Mq { get; set; } - private BlockingStaticTaskScheduler FileTaskScheduler { get; set; } - private BlockingStaticTaskScheduler TreeTaskScheduler { get; set; } + private FileTaskScheduler FileTaskScheduler { get; set; } + private TreeTaskScheduler TreeTaskScheduler { get; set; } private FileScanner FileScanner { get; set; } public TreeWalker() @@ -38,18 +38,19 @@ public void WalkTree(string currentDir) // check if we actually like the files foreach (string file in files) { - FileTaskScheduler.New(() => - { - try - { - FileScanner.ScanFile(file); - } - catch (Exception e) - { - Mq.Error("Exception in FileScanner task for file " + file); - Mq.Trace(e.ToString()); - } - }); + FileTaskScheduler.New(FileScanner.ScanFile, file); + //FileTaskScheduler.New(() => + //{ + // try + // { + // FileScanner.ScanFile(file); + // } + // catch (Exception e) + // { + // Mq.Error("Exception in FileScanner task for file " + file); + // Mq.Trace(e.ToString()); + // } + //}); } } catch (UnauthorizedAccessException) @@ -104,18 +105,19 @@ public void WalkTree(string currentDir) } if (scanDir == true) { - TreeTaskScheduler.New(() => - { - try - { - WalkTree(dirStr); - } - catch (Exception e) - { - Mq.Error("Exception in TreeWalker task for dir " + dirStr); - Mq.Error(e.ToString()); - } - }); + TreeTaskScheduler.New(WalkTree, dirStr); + //TreeTaskScheduler.New(() => + //{ + // try + // { + // WalkTree(dirStr); + // } + // catch (Exception e) + // { + // Mq.Error("Exception in TreeWalker task for dir " + dirStr); + // Mq.Error(e.ToString()); + // } + //}); } else { diff --git a/Snaffler/Config.cs b/Snaffler/Config.cs index 0a7e9a57..0ec0ddd0 100644 --- a/Snaffler/Config.cs +++ b/Snaffler/Config.cs @@ -106,6 +106,8 @@ private static Options ParseImpl(string[] args) ValueArgument logType = new ValueArgument('t', "logtype", "Type of log you would like to output. Currently supported options are plain and JSON. Defaults to plain."); ValueArgument timeOutArg = new ValueArgument('e', "timeout", "Interval between status updates (in minutes) also acts as a timeout for AD data to be gathered via LDAP. Turn this knob up if you aren't getting any computers from AD when you run Snaffler through a proxy or other slow link. Default = 5"); + ValueArgument taskFile = new ValueArgument('1', "taskfile", "Save tasks as they are created to a file to allow for resuming mid-operation."); + ValueArgument resumeFrom = new ValueArgument('2', "resumefrom", "Resume tasks from a file generated with --taskfile."); // list of letters i haven't used yet: gnqw CommandLineParser.CommandLineParser parser = new CommandLineParser.CommandLineParser(); @@ -132,6 +134,8 @@ private static Options ParseImpl(string[] args) parser.Arguments.Add(ruleDirArg); parser.Arguments.Add(logType); parser.Arguments.Add(compExclusionArg); + parser.Arguments.Add(taskFile); + parser.Arguments.Add(resumeFrom); // extra check to handle builtin behaviour from cmd line arg parser if ((args.Contains("--help") || args.Contains("/?") || args.Contains("help") || args.Contains("-h") || args.Length == 0)) @@ -150,6 +154,16 @@ private static Options ParseImpl(string[] args) { parser.ParseCommandLine(args); + if (taskFile.Parsed) + { + parsedConfig.TaskFile = taskFile.Value; + } + + if (resumeFrom.Parsed) + { + parsedConfig.ResumeFrom = resumeFrom.Value; + } + if (timeOutArg.Parsed && !String.IsNullOrWhiteSpace(timeOutArg.Value)) { int timeOutVal; From d3e72020b95b7c22961a0f3dd008f215cf31fe05 Mon Sep 17 00:00:00 2001 From: p0rtL6 Date: Mon, 7 Apr 2025 18:18:06 -0700 Subject: [PATCH 2/3] Bug Fixes --- .../Concurrency/BlockingTaskScheduler.cs | 119 +++++++++------ SnaffCore/ShareFind/ShareFinder.cs | 15 +- SnaffCore/SnaffCon.cs | 139 ++++++++---------- SnaffCore/TreeWalk/TreeWalker.cs | 24 --- 4 files changed, 140 insertions(+), 157 deletions(-) diff --git a/SnaffCore/Concurrency/BlockingTaskScheduler.cs b/SnaffCore/Concurrency/BlockingTaskScheduler.cs index 876948f0..e9217be4 100644 --- a/SnaffCore/Concurrency/BlockingTaskScheduler.cs +++ b/SnaffCore/Concurrency/BlockingTaskScheduler.cs @@ -73,21 +73,50 @@ public void New(Action action) } } - public enum TaskType + public enum TaskFileType { None = 0, Share = 1, Tree = 2, File = 3 } + + public enum TaskFileEntryStatus + { + Pending = 0, + Completed = 1, + } public struct TaskFileEntry { public TaskFileEntryStatus status; public string guid; - public TaskType type; + public TaskFileType type; public string input; + public override string ToString() + { + StringBuilder stringBuilder = new StringBuilder(); + + stringBuilder.Append(status.ToString()); + stringBuilder.Append("|"); + stringBuilder.Append(guid); + stringBuilder.Append("|"); + stringBuilder.Append(type.ToString()); + stringBuilder.Append("|"); + stringBuilder.Append(input); + + return stringBuilder.ToString(); + } + + public TaskFileEntry(TaskFileType type, string input) + { + guid = Guid.NewGuid().ToString(); + status = TaskFileEntryStatus.Pending; + this.type = type; + this.input = input; + } + public TaskFileEntry(string entryLine) { string[] lineParts = entryLine.Split('|'); @@ -95,25 +124,11 @@ public TaskFileEntry(string entryLine) status = (TaskFileEntryStatus)Enum.Parse(typeof(TaskFileEntryStatus), lineParts[0]); guid = lineParts[1]; - if (status == TaskFileEntryStatus.Pending) - { - type = (TaskType)Enum.Parse(typeof(TaskType), lineParts[2]); - input = lineParts[3]; - } - else - { - type = TaskType.None; - input = null; - } + type = (TaskFileType)Enum.Parse(typeof(TaskFileType), lineParts[2]); + input = lineParts[3]; } } - public enum TaskFileEntryStatus - { - Pending = 0, - Completed = 1, - } - public class ResumingTaskScheduler : BlockingStaticTaskScheduler { private static readonly object WriteLock = new object(); @@ -130,7 +145,7 @@ public ResumingTaskScheduler(int threads, int maxBacklog) : base(threads, maxBac public static void SetAlreadyHandledTasks(TaskFileEntry[] taskFileEntries) { - AlreadyHandledTasks = taskFileEntries.Select(entry => entry.input).ToArray(); + AlreadyHandledTasks = taskFileEntries.Select(entry => entry.input).Distinct().ToArray(); } public static bool IsTaskAlreadyHandled(string input) @@ -139,11 +154,16 @@ public static bool IsTaskAlreadyHandled(string input) return AlreadyHandledTasks.Contains(input); } - public void New(TaskType taskType, Action action, string input) + public void New(TaskFileType taskType, Action action, string input) + { + New(taskType, action, input, false); + } + + public void New(TaskFileType taskType, Action action, string input, bool ignoreAlreadyHandled) { - if (IsTaskAlreadyHandled(input)) return; + if (!ignoreAlreadyHandled && IsTaskAlreadyHandled(input)) return; - string taskGuid = SaveTask(taskType, input); + TaskFileEntry? taskFileEntry = SaveTask(taskType, input); New(() => { @@ -157,7 +177,7 @@ public void New(TaskType taskType, Action action, string input) Mq.Error(e.ToString()); } - CompleteTask(taskGuid); + CompleteTask(taskFileEntry); }); } @@ -177,38 +197,34 @@ public static void CloseTaskFile() } } - internal string SaveTask(TaskType taskType, string input) + internal TaskFileEntry? SaveTask(TaskFileType taskType, string input) { // task file is not set, we are not saving tasks if (fileWriter == null) return null; - string taskGuid = Guid.NewGuid().ToString(); - - StringBuilder taskEntryBuilder = new StringBuilder(); - - taskEntryBuilder.Append("Pending|"); - taskEntryBuilder.Append(taskGuid); - taskEntryBuilder.Append("|"); - taskEntryBuilder.Append(taskType.ToString()); - taskEntryBuilder.Append("|"); - taskEntryBuilder.Append(input); + TaskFileEntry taskFileEntry = new TaskFileEntry(taskType, input); lock (WriteLock) { - fileWriter.WriteLine(taskEntryBuilder.ToString()); + fileWriter.WriteLine(taskFileEntry.ToString()); fileWriter.Flush(); } - return taskGuid; + return taskFileEntry; } - internal void CompleteTask(string taskGuid) + public static void CompleteTask(TaskFileEntry? taskFileEntry) { if (fileWriter == null) return; + if (!taskFileEntry.HasValue) return; + + TaskFileEntry taskFileEntryValue = taskFileEntry.Value; + + taskFileEntryValue.status = TaskFileEntryStatus.Completed; lock (WriteLock) { - fileWriter.WriteLine("Completed|" + taskGuid); + fileWriter.WriteLine(taskFileEntryValue.ToString()); fileWriter.Flush(); } } @@ -218,9 +234,14 @@ public class ShareTaskScheduler : ResumingTaskScheduler { public ShareTaskScheduler(int threads, int maxBacklog) : base(threads, maxBacklog) { } - public void New(Action action, string file) + public void New(Action action, string share) { - New(TaskType.Share, action, file); + New(action, share, false); + } + + public void New(Action action, string share, bool ignoreAlreadyHandled) + { + New(TaskFileType.Share, action, share, ignoreAlreadyHandled); } } @@ -228,9 +249,14 @@ public class TreeTaskScheduler : ResumingTaskScheduler { public TreeTaskScheduler(int threads, int maxBacklog) : base(threads, maxBacklog) { } - public void New(Action action, string file) + public void New(Action action, string tree) { - New(TaskType.Tree, action, file); + New(action, tree, false); + } + + public void New(Action action, string tree, bool ignoreAlreadyHandled) + { + New(TaskFileType.Tree, action, tree, ignoreAlreadyHandled); } } @@ -238,9 +264,14 @@ public class FileTaskScheduler : ResumingTaskScheduler { public FileTaskScheduler(int threads, int maxBacklog) : base(threads, maxBacklog) { } - public void New(Action action, string file) + public void New(Action action, string tree) + { + New(action, tree, false); + } + + public void New(Action action, string file, bool ignoreAlreadyHandled) { - New(TaskType.File, action, file); + New(TaskFileType.File, action, file, ignoreAlreadyHandled); } } diff --git a/SnaffCore/ShareFind/ShareFinder.cs b/SnaffCore/ShareFind/ShareFinder.cs index 4d1395c1..6bcb9543 100644 --- a/SnaffCore/ShareFind/ShareFinder.cs +++ b/SnaffCore/ShareFind/ShareFinder.cs @@ -17,7 +17,7 @@ namespace SnaffCore.ShareFind public class ShareFinder { private BlockingMq Mq { get; set; } - private BlockingStaticTaskScheduler TreeTaskScheduler { get; set; } + private TreeTaskScheduler TreeTaskScheduler { get; set; } private TreeWalker TreeWalker { get; set; } //private EffectivePermissions effectivePermissions { get; set; } = new EffectivePermissions(MyOptions.CurrentUser); @@ -184,18 +184,7 @@ internal void GetComputerShares(string computer) if (MyOptions.ScanFoundShares) { Mq.Trace("Creating a TreeWalker task for " + shareResult.SharePath); - TreeTaskScheduler.New(() => - { - try - { - TreeWalker.WalkTree(shareResult.SharePath); - } - catch (Exception e) - { - Mq.Error("Exception in TreeWalker task for share " + shareResult.SharePath); - Mq.Error(e.ToString()); - } - }); + TreeTaskScheduler.New(TreeWalker.WalkTree, shareResult.SharePath); } Mq.ShareResult(shareResult); } diff --git a/SnaffCore/SnaffCon.cs b/SnaffCore/SnaffCon.cs index 092c48f3..cea5f80d 100644 --- a/SnaffCore/SnaffCon.cs +++ b/SnaffCore/SnaffCon.cs @@ -100,87 +100,99 @@ public void Execute() if (MyOptions.ResumeFrom != null && MyOptions.ResumeFrom.Length > 0) { - // parse task file into structured data + // Parse task file into structured data TaskFileEntry[] taskFileEntries = File.ReadAllLines(MyOptions.ResumeFrom).Select(line => new TaskFileEntry(line)).ToArray(); + // Set the internal list on the scheduler so it can skip tasks (adds pending too because those will be handled below) + ResumingTaskScheduler.SetAlreadyHandledTasks(taskFileEntries); + + // Parse the completed and pending tasks for dispatching TaskFileEntry[] completedTasks = taskFileEntries.Where(entry => entry.status == TaskFileEntryStatus.Completed).ToArray(); + + // If the user has specified a new file to save to, append completed tasks for continuity + foreach (TaskFileEntry completedTask in completedTasks) + { + ResumingTaskScheduler.CompleteTask(completedTask); + } + TaskFileEntry[] pendingTasks = taskFileEntries.Where(entry => (entry.status == TaskFileEntryStatus.Pending && !completedTasks.Any(completedEntry => completedEntry.guid == entry.guid))).ToArray(); + // Dispatch pending tasks foreach (TaskFileEntry pendingFileTask in pendingTasks) { switch (pendingFileTask.type) { - case TaskType.Share: + case TaskFileType.Share: ShareFinder shareFinder = new ShareFinder(); - ShareTaskScheduler.New(shareFinder.GetComputerShares, pendingFileTask.input); + ShareTaskScheduler.New(shareFinder.GetComputerShares, pendingFileTask.input, true); break; - case TaskType.Tree: - TreeTaskScheduler.New(TreeWalker.WalkTree, pendingFileTask.input); + case TaskFileType.Tree: + TreeTaskScheduler.New(TreeWalker.WalkTree, pendingFileTask.input, true); break; - case TaskType.File: - FileTaskScheduler.New(FileScanner.ScanFile, pendingFileTask.input); + case TaskFileType.File: + FileTaskScheduler.New(FileScanner.ScanFile, pendingFileTask.input, true); break; } } } - else + + // Continue with normal operation, anything handled from the file above will be skipped internally by the scheduler + + // If we want to hunt for user IDs, we need data from the running user's domain. + // Future - walk trusts + if (MyOptions.DomainUserRules) { - // If we want to hunt for user IDs, we need data from the running user's domain. - // Future - walk trusts - if (MyOptions.DomainUserRules) - { - DomainUserDiscovery(); - } + DomainUserDiscovery(); + } - // Explicit folder setting overrides DFS - if (MyOptions.PathTargets.Count != 0 && (MyOptions.DfsShareDiscovery || MyOptions.DfsOnly)) + // Explicit folder setting overrides DFS + if (MyOptions.PathTargets.Count != 0 && (MyOptions.DfsShareDiscovery || MyOptions.DfsOnly)) + { + DomainDfsDiscovery(); + } + + if (MyOptions.PathTargets.Count == 0 && MyOptions.ComputerTargets == null) + { + if (MyOptions.DfsSharesDict.Count == 0) { + Mq.Info("Invoking DFS Discovery because no ComputerTargets or PathTargets were specified"); DomainDfsDiscovery(); } - if (MyOptions.PathTargets.Count == 0 && MyOptions.ComputerTargets == null) + if (!MyOptions.DfsOnly) { - if (MyOptions.DfsSharesDict.Count == 0) - { - Mq.Info("Invoking DFS Discovery because no ComputerTargets or PathTargets were specified"); - DomainDfsDiscovery(); - } - - if (!MyOptions.DfsOnly) - { - Mq.Info("Invoking full domain computer discovery."); - DomainTargetDiscovery(); - } - else + Mq.Info("Invoking full domain computer discovery."); + DomainTargetDiscovery(); + } + else + { + Mq.Info("Skipping domain computer discovery."); + foreach (string share in MyOptions.DfsSharesDict.Keys) { - Mq.Info("Skipping domain computer discovery."); - foreach (string share in MyOptions.DfsSharesDict.Keys) + if (!MyOptions.PathTargets.Contains(share)) { - if (!MyOptions.PathTargets.Contains(share)) - { - MyOptions.PathTargets.Add(share); - } + MyOptions.PathTargets.Add(share); } - Mq.Info("Starting TreeWalker tasks on DFS shares."); - FileDiscovery(MyOptions.PathTargets.ToArray()); } - } - // otherwise we should have a set of path targets... - else if (MyOptions.PathTargets.Count != 0) - { + Mq.Info("Starting TreeWalker tasks on DFS shares."); FileDiscovery(MyOptions.PathTargets.ToArray()); } - // or we've been told what computers to hit... - else if (MyOptions.ComputerTargets != null) - { - ShareDiscovery(MyOptions.ComputerTargets); - } + } + // otherwise we should have a set of path targets... + else if (MyOptions.PathTargets.Count != 0) + { + FileDiscovery(MyOptions.PathTargets.ToArray()); + } + // or we've been told what computers to hit... + else if (MyOptions.ComputerTargets != null) + { + ShareDiscovery(MyOptions.ComputerTargets); + } - // but if that hasn't been done, something has gone wrong. - else - { - Mq.Error("OctoParrot says: AWK! I SHOULDN'T BE!"); - } + // but if that hasn't been done, something has gone wrong. + else + { + Mq.Error("OctoParrot says: AWK! I SHOULDN'T BE!"); } waitHandle.WaitOne(); @@ -336,19 +348,6 @@ private void ShareDiscovery(string[] computerTargets) Mq.Trace("Creating a ShareFinder task for " + computer); ShareFinder shareFinder = new ShareFinder(); ShareTaskScheduler.New(shareFinder.GetComputerShares, computer); - //ShareTaskScheduler.New(() => - //{ - // try - // { - // ShareFinder shareFinder = new ShareFinder(); - // shareFinder.GetComputerShares(computer); - // } - // catch (Exception e) - // { - // Mq.Error("Exception in ShareFinder task for host " + computer); - // Mq.Error(e.ToString()); - // } - //}); } Mq.Info("Created all sharefinder tasks."); } @@ -411,18 +410,6 @@ private void FileDiscovery(string[] pathTargets) // TreeWalker Task Creation - this kicks off the rest of the flow Mq.Info("Creating a TreeWalker task for " + pathTarget); TreeTaskScheduler.New(TreeWalker.WalkTree, pathTarget); - //TreeTaskScheduler.New(() => - //{ - // try - // { - // TreeWalker.WalkTree(pathTarget); - // } - // catch (Exception e) - // { - // Mq.Error("Exception in TreeWalker task for path " + pathTarget); - // Mq.Error(e.ToString()); - // } - //}); } Mq.Info("Created all TreeWalker tasks."); diff --git a/SnaffCore/TreeWalk/TreeWalker.cs b/SnaffCore/TreeWalk/TreeWalker.cs index 825c90f3..af6f89d5 100644 --- a/SnaffCore/TreeWalk/TreeWalker.cs +++ b/SnaffCore/TreeWalk/TreeWalker.cs @@ -39,18 +39,6 @@ public void WalkTree(string currentDir) foreach (string file in files) { FileTaskScheduler.New(FileScanner.ScanFile, file); - //FileTaskScheduler.New(() => - //{ - // try - // { - // FileScanner.ScanFile(file); - // } - // catch (Exception e) - // { - // Mq.Error("Exception in FileScanner task for file " + file); - // Mq.Trace(e.ToString()); - // } - //}); } } catch (UnauthorizedAccessException) @@ -106,18 +94,6 @@ public void WalkTree(string currentDir) if (scanDir == true) { TreeTaskScheduler.New(WalkTree, dirStr); - //TreeTaskScheduler.New(() => - //{ - // try - // { - // WalkTree(dirStr); - // } - // catch (Exception e) - // { - // Mq.Error("Exception in TreeWalker task for dir " + dirStr); - // Mq.Error(e.ToString()); - // } - //}); } else { From e081d8d11697b9936003c21ebd8a1d13fb9ba439 Mon Sep 17 00:00:00 2001 From: p0rtL6 Date: Tue, 8 Apr 2025 17:06:15 -0700 Subject: [PATCH 3/3] Update resume feature to save at defined intervals --- .../Concurrency/BlockingTaskScheduler.cs | 122 +++++---------- SnaffCore/Config/Options.cs | 1 + SnaffCore/SnaffCon.cs | 144 +++++++++--------- Snaffler/Config.cs | 16 ++ 4 files changed, 133 insertions(+), 150 deletions(-) diff --git a/SnaffCore/Concurrency/BlockingTaskScheduler.cs b/SnaffCore/Concurrency/BlockingTaskScheduler.cs index e9217be4..47d847cb 100644 --- a/SnaffCore/Concurrency/BlockingTaskScheduler.cs +++ b/SnaffCore/Concurrency/BlockingTaskScheduler.cs @@ -6,6 +6,8 @@ using System.Text; using System.Threading; using System.Threading.Tasks; +using System.Timers; +using static SnaffCore.Config.Options; namespace SnaffCore.Concurrency { @@ -131,10 +133,9 @@ public TaskFileEntry(string entryLine) public class ResumingTaskScheduler : BlockingStaticTaskScheduler { - private static readonly object WriteLock = new object(); - private static StreamWriter fileWriter; - - private static string[] AlreadyHandledTasks; + private static readonly Dictionary> pendingTasks = new Dictionary>(); + private static readonly object writeLock = new object(); + private static int pendingSaveCalls = 0; internal BlockingMq Mq { get; } @@ -143,89 +144,65 @@ public ResumingTaskScheduler(int threads, int maxBacklog) : base(threads, maxBac this.Mq = BlockingMq.GetMq(); } - public static void SetAlreadyHandledTasks(TaskFileEntry[] taskFileEntries) - { - AlreadyHandledTasks = taskFileEntries.Select(entry => entry.input).Distinct().ToArray(); - } - - public static bool IsTaskAlreadyHandled(string input) - { - if (AlreadyHandledTasks == null) return false; - return AlreadyHandledTasks.Contains(input); - } - - public void New(TaskFileType taskType, Action action, string input) - { - New(taskType, action, input, false); - } - - public void New(TaskFileType taskType, Action action, string input, bool ignoreAlreadyHandled) + public void New(string taskType, Action action, string path) { - if (!ignoreAlreadyHandled && IsTaskAlreadyHandled(input)) return; + string guid = null; - TaskFileEntry? taskFileEntry = SaveTask(taskType, input); + if (MyOptions.TaskFile != null) + { + guid = Guid.NewGuid().ToString(); + pendingTasks.Add(guid, new Tuple(taskType, path)); + } New(() => { try { - action(input); + action(path); } catch (Exception e) { - Mq.Error("Exception in " + taskType.ToString() + " task for host " + input); + Mq.Error("Exception in " + taskType.ToString() + " task for host " + path); Mq.Error(e.ToString()); } - CompleteTask(taskFileEntry); + if (guid != null) pendingTasks.Remove(guid); }); } - public static void SetTaskFile(string path) + public static void SaveState(object sender, ElapsedEventArgs e) { - fileWriter = new StreamWriter(path); + SaveState(); } - public static void CloseTaskFile() + public static void SaveState() { - if (fileWriter != null) { - lock (WriteLock) - { - fileWriter.Close(); - fileWriter = null; - } - } - } - - internal TaskFileEntry? SaveTask(TaskFileType taskType, string input) - { - // task file is not set, we are not saving tasks - if (fileWriter == null) return null; + // Guard against the possibility that someone forgot to check this + if (MyOptions.TaskFile == null) return; - TaskFileEntry taskFileEntry = new TaskFileEntry(taskType, input); + // This blocks more than one save call from being buffered at a time + // Prevents a situation where a bunch of buffered calls wait for the lock + // But still allows for you to "write continously" if you set an interval shorter than the file write time + if (pendingSaveCalls > 1) return; + pendingSaveCalls++; - lock (WriteLock) + // In case the file takes longer to write than the set interval + lock (writeLock) { - fileWriter.WriteLine(taskFileEntry.ToString()); - fileWriter.Flush(); - } - - return taskFileEntry; - } - - public static void CompleteTask(TaskFileEntry? taskFileEntry) - { - if (fileWriter == null) return; - if (!taskFileEntry.HasValue) return; + using (StreamWriter fileWriter = new StreamWriter(MyOptions.TaskFile, false)) + { + // Copy the values into the array to avoid an error in case the pending tasks are changed during the write loop + Tuple[] valuesSnapshot = pendingTasks.Values.ToArray(); - TaskFileEntry taskFileEntryValue = taskFileEntry.Value; + foreach (Tuple value in valuesSnapshot) + { + fileWriter.WriteLine($"{value.Item1}|{value.Item2}"); + } - taskFileEntryValue.status = TaskFileEntryStatus.Completed; + fileWriter.Flush(); + } - lock (WriteLock) - { - fileWriter.WriteLine(taskFileEntryValue.ToString()); - fileWriter.Flush(); + pendingSaveCalls--; } } } @@ -236,12 +213,7 @@ public ShareTaskScheduler(int threads, int maxBacklog) : base(threads, maxBacklo public void New(Action action, string share) { - New(action, share, false); - } - - public void New(Action action, string share, bool ignoreAlreadyHandled) - { - New(TaskFileType.Share, action, share, ignoreAlreadyHandled); + New("share", action, share); } } @@ -251,12 +223,7 @@ public TreeTaskScheduler(int threads, int maxBacklog) : base(threads, maxBacklog public void New(Action action, string tree) { - New(action, tree, false); - } - - public void New(Action action, string tree, bool ignoreAlreadyHandled) - { - New(TaskFileType.Tree, action, tree, ignoreAlreadyHandled); + New("tree", action, tree); } } @@ -264,14 +231,9 @@ public class FileTaskScheduler : ResumingTaskScheduler { public FileTaskScheduler(int threads, int maxBacklog) : base(threads, maxBacklog) { } - public void New(Action action, string tree) - { - New(action, tree, false); - } - - public void New(Action action, string file, bool ignoreAlreadyHandled) + public void New(Action action, string file) { - New(TaskFileType.File, action, file, ignoreAlreadyHandled); + New("file", action, file); } } diff --git a/SnaffCore/Config/Options.cs b/SnaffCore/Config/Options.cs index 4c6c6eec..44761572 100644 --- a/SnaffCore/Config/Options.cs +++ b/SnaffCore/Config/Options.cs @@ -16,6 +16,7 @@ public partial class Options // Pause and resume functionality public string TaskFile { get; set; } + public double TaskFileTimeOut { get; set; } = 5; public string ResumeFrom { get; set; } // Manual Targeting Options diff --git a/SnaffCore/SnaffCon.cs b/SnaffCore/SnaffCon.cs index cea5f80d..9094454d 100644 --- a/SnaffCore/SnaffCon.cs +++ b/SnaffCore/SnaffCon.cs @@ -93,110 +93,114 @@ public void Execute() statusUpdateTimer.Elapsed += TimedStatusUpdate; statusUpdateTimer.Start(); - if (MyOptions.TaskFile != null && MyOptions.TaskFile.Length > 0) + + if (MyOptions.TaskFile != null) { - ResumingTaskScheduler.SetTaskFile(MyOptions.TaskFile); + Timer saveStateTimer = new Timer(TimeSpan.FromMinutes(MyOptions.TaskFileTimeOut).TotalMilliseconds) { AutoReset = true }; + saveStateTimer.Elapsed += ResumingTaskScheduler.SaveState; + saveStateTimer.Start(); } - if (MyOptions.ResumeFrom != null && MyOptions.ResumeFrom.Length > 0) + if (MyOptions.ResumeFrom != null) { - // Parse task file into structured data - TaskFileEntry[] taskFileEntries = File.ReadAllLines(MyOptions.ResumeFrom).Select(line => new TaskFileEntry(line)).ToArray(); + // Read the task file and parse it into tuples + Tuple[] taskFileEntries = File.ReadLines(MyOptions.ResumeFrom).Select(line => + { + string[] parts = line.Split('|'); + return parts.Length == 2 ? new Tuple(parts[0], parts[1]) : null; + }).Where(tuple => tuple != null).ToArray(); - // Set the internal list on the scheduler so it can skip tasks (adds pending too because those will be handled below) - ResumingTaskScheduler.SetAlreadyHandledTasks(taskFileEntries); + // Get all shares, they are non-recursive so just save all of them + Tuple[] shareEntries = taskFileEntries.Where(entry => entry.Item1 == "share").ToArray(); - // Parse the completed and pending tasks for dispatching - TaskFileEntry[] completedTasks = taskFileEntries.Where(entry => entry.status == TaskFileEntryStatus.Completed).ToArray(); - - // If the user has specified a new file to save to, append completed tasks for continuity - foreach (TaskFileEntry completedTask in completedTasks) - { - ResumingTaskScheduler.CompleteTask(completedTask); - } + // Remove all entries where the path starts with a pending share + taskFileEntries = taskFileEntries.Where(entry => !shareEntries.Any(shareEntry => entry.Item2.StartsWith("\\\\" + shareEntry.Item2 + "\\"))).ToArray(); + + // Get all tree entires, they are recursive so we need to find the shortest path for each shared base + Tuple[] treeEntries = taskFileEntries.Where(entry => entry.Item1 == "tree" && !taskFileEntries.Any(otherEntry => entry.Item2.StartsWith(otherEntry.Item2 + "\\"))).ToArray(); - TaskFileEntry[] pendingTasks = taskFileEntries.Where(entry => (entry.status == TaskFileEntryStatus.Pending && !completedTasks.Any(completedEntry => completedEntry.guid == entry.guid))).ToArray(); + // Remove all entries where the path starts with one of the base pending tree + taskFileEntries = taskFileEntries.Where(entry => !treeEntries.Any(treeEntry => entry.Item2.StartsWith(treeEntry.Item2 + "\\"))).ToArray(); - // Dispatch pending tasks - foreach (TaskFileEntry pendingFileTask in pendingTasks) + // Tasks should be deduplicated now, dispatch what is left pending + foreach (Tuple entry in taskFileEntries) { - switch (pendingFileTask.type) + switch (entry.Item1) { - case TaskFileType.Share: + case "share": ShareFinder shareFinder = new ShareFinder(); - ShareTaskScheduler.New(shareFinder.GetComputerShares, pendingFileTask.input, true); + ShareTaskScheduler.New(shareFinder.GetComputerShares, entry.Item2); break; - case TaskFileType.Tree: - TreeTaskScheduler.New(TreeWalker.WalkTree, pendingFileTask.input, true); + case "tree": + TreeTaskScheduler.New(TreeWalker.WalkTree, entry.Item2); break; - case TaskFileType.File: - FileTaskScheduler.New(FileScanner.ScanFile, pendingFileTask.input, true); + case "file": + FileTaskScheduler.New(FileScanner.ScanFile, entry.Item2); break; } } } - - // Continue with normal operation, anything handled from the file above will be skipped internally by the scheduler - - // If we want to hunt for user IDs, we need data from the running user's domain. - // Future - walk trusts - if (MyOptions.DomainUserRules) - { - DomainUserDiscovery(); - } - - // Explicit folder setting overrides DFS - if (MyOptions.PathTargets.Count != 0 && (MyOptions.DfsShareDiscovery || MyOptions.DfsOnly)) - { - DomainDfsDiscovery(); - } - - if (MyOptions.PathTargets.Count == 0 && MyOptions.ComputerTargets == null) + else { - if (MyOptions.DfsSharesDict.Count == 0) + // If we want to hunt for user IDs, we need data from the running user's domain. + // Future - walk trusts + if (MyOptions.DomainUserRules) { - Mq.Info("Invoking DFS Discovery because no ComputerTargets or PathTargets were specified"); - DomainDfsDiscovery(); + DomainUserDiscovery(); } - if (!MyOptions.DfsOnly) + // Explicit folder setting overrides DFS + if (MyOptions.PathTargets.Count != 0 && (MyOptions.DfsShareDiscovery || MyOptions.DfsOnly)) { - Mq.Info("Invoking full domain computer discovery."); - DomainTargetDiscovery(); + DomainDfsDiscovery(); } - else + + if (MyOptions.PathTargets.Count == 0 && MyOptions.ComputerTargets == null) { - Mq.Info("Skipping domain computer discovery."); - foreach (string share in MyOptions.DfsSharesDict.Keys) + if (MyOptions.DfsSharesDict.Count == 0) { - if (!MyOptions.PathTargets.Contains(share)) + Mq.Info("Invoking DFS Discovery because no ComputerTargets or PathTargets were specified"); + DomainDfsDiscovery(); + } + + if (!MyOptions.DfsOnly) + { + Mq.Info("Invoking full domain computer discovery."); + DomainTargetDiscovery(); + } + else + { + Mq.Info("Skipping domain computer discovery."); + foreach (string share in MyOptions.DfsSharesDict.Keys) { - MyOptions.PathTargets.Add(share); + if (!MyOptions.PathTargets.Contains(share)) + { + MyOptions.PathTargets.Add(share); + } } + Mq.Info("Starting TreeWalker tasks on DFS shares."); + FileDiscovery(MyOptions.PathTargets.ToArray()); } - Mq.Info("Starting TreeWalker tasks on DFS shares."); + } + // otherwise we should have a set of path targets... + else if (MyOptions.PathTargets.Count != 0) + { FileDiscovery(MyOptions.PathTargets.ToArray()); } - } - // otherwise we should have a set of path targets... - else if (MyOptions.PathTargets.Count != 0) - { - FileDiscovery(MyOptions.PathTargets.ToArray()); - } - // or we've been told what computers to hit... - else if (MyOptions.ComputerTargets != null) - { - ShareDiscovery(MyOptions.ComputerTargets); - } + // or we've been told what computers to hit... + else if (MyOptions.ComputerTargets != null) + { + ShareDiscovery(MyOptions.ComputerTargets); + } - // but if that hasn't been done, something has gone wrong. - else - { - Mq.Error("OctoParrot says: AWK! I SHOULDN'T BE!"); + // but if that hasn't been done, something has gone wrong. + else + { + Mq.Error("OctoParrot says: AWK! I SHOULDN'T BE!"); + } } waitHandle.WaitOne(); - ResumingTaskScheduler.CloseTaskFile(); StatusUpdate(); diff --git a/Snaffler/Config.cs b/Snaffler/Config.cs index 0ec0ddd0..ce66e366 100644 --- a/Snaffler/Config.cs +++ b/Snaffler/Config.cs @@ -108,6 +108,7 @@ private static Options ParseImpl(string[] args) "Interval between status updates (in minutes) also acts as a timeout for AD data to be gathered via LDAP. Turn this knob up if you aren't getting any computers from AD when you run Snaffler through a proxy or other slow link. Default = 5"); ValueArgument taskFile = new ValueArgument('1', "taskfile", "Save tasks as they are created to a file to allow for resuming mid-operation."); ValueArgument resumeFrom = new ValueArgument('2', "resumefrom", "Resume tasks from a file generated with --taskfile."); + ValueArgument taskFileTimeOut = new ValueArgument('3', "taskfiletimeout", "Interval between saving tasks to the task file (in minutes). Default = 5"); // list of letters i haven't used yet: gnqw CommandLineParser.CommandLineParser parser = new CommandLineParser.CommandLineParser(); @@ -135,6 +136,7 @@ private static Options ParseImpl(string[] args) parser.Arguments.Add(logType); parser.Arguments.Add(compExclusionArg); parser.Arguments.Add(taskFile); + parser.Arguments.Add(taskFileTimeOut); parser.Arguments.Add(resumeFrom); // extra check to handle builtin behaviour from cmd line arg parser @@ -159,6 +161,20 @@ private static Options ParseImpl(string[] args) parsedConfig.TaskFile = taskFile.Value; } + if (taskFileTimeOut.Parsed && !String.IsNullOrWhiteSpace(taskFileTimeOut.Value)) + { + double timeOutVal; + if (double.TryParse(taskFileTimeOut.Value, out timeOutVal)) + { + Mq.Info("Set task file saving interval to " + timeOutVal.ToString() + " minutes."); + parsedConfig.TaskFileTimeOut = timeOutVal; + } + else + { + Mq.Error("Invalid task file timeout value passed, defaulting to 5 mins."); + } + } + if (resumeFrom.Parsed) { parsedConfig.ResumeFrom = resumeFrom.Value;