diff --git a/SnaffCore/Checkpoint/ScanState.cs b/SnaffCore/Checkpoint/ScanState.cs new file mode 100644 index 00000000..825b9e8b --- /dev/null +++ b/SnaffCore/Checkpoint/ScanState.cs @@ -0,0 +1,310 @@ +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.IO; +using System.Text; +using SnaffCore.Concurrency; + +namespace SnaffCore.Checkpoint +{ + /// + /// Tracks scan progress so a run can be resumed after an interruption + /// (network drop, laptop sleep, crash, Ctrl-C, etc). + /// + /// Design goals: + /// - No external dependencies (works fine on .NET Framework 4.5.1). + /// - Crash-safe: every completed unit of work is appended to disk + /// immediately and flushed, instead of being buffered and saved + /// periodically. If the machine dies mid-scan, we only lose the + /// (small) amount of work that was in flight at that instant - + /// nothing that was already marked done. + /// - Same command line every time: point -w/--checkpoint at a file. + /// First run creates it. Any later run against the same file + /// skips everything already recorded there and carries on. + /// + /// File format is a simple pipe-delimited append log: + /// D| -> directory fully enumerated + /// C| -> share list for computer fully enumerated + /// S|| -> a share that was found scannable on that computer + /// + public class ScanState + { + private static readonly object _initLock = new object(); + private static ScanState _instance; + + public static ScanState Instance + { + get + { + if (_instance == null) + { + lock (_initLock) + { + if (_instance == null) + { + _instance = new ScanState(); + } + } + } + return _instance; + } + } + + private readonly object _writeLock = new object(); + private ConcurrentDictionary _completedDirs; + private ConcurrentDictionary _completedComputers; + private ConcurrentDictionary> _computerShares; + private ConcurrentDictionary _pendingFileCounts; + + private string _checkpointPath; + private StreamWriter _writer; + public bool Enabled { get; private set; } + + private ScanState() + { + _completedDirs = new ConcurrentDictionary(StringComparer.OrdinalIgnoreCase); + _completedComputers = new ConcurrentDictionary(StringComparer.OrdinalIgnoreCase); + _computerShares = new ConcurrentDictionary>(StringComparer.OrdinalIgnoreCase); + _pendingFileCounts = new ConcurrentDictionary(StringComparer.OrdinalIgnoreCase); + } + + /// + /// Call once at startup. If the checkpoint file already has content, + /// it's replayed into memory so this run resumes instead of restarting. + /// + public void Load(string checkpointPath) + { + if (string.IsNullOrWhiteSpace(checkpointPath)) + { + Enabled = false; + return; + } + + BlockingMq mq = BlockingMq.GetMq(); + _checkpointPath = checkpointPath; + Enabled = true; + + int resumedDirs = 0; + int resumedComputers = 0; + int resumedShares = 0; + + if (File.Exists(_checkpointPath)) + { + using (FileStream fs = new FileStream(_checkpointPath, FileMode.Open, FileAccess.Read, FileShare.ReadWrite)) + using (StreamReader reader = new StreamReader(fs, Encoding.UTF8)) + { + string line; + while ((line = reader.ReadLine()) != null) + { + if (string.IsNullOrWhiteSpace(line)) + { + continue; + } + + string[] parts = line.Split(new[] { '|' }, 3); + if (parts.Length < 2) + { + continue; + } + + switch (parts[0]) + { + case "D": + if (_completedDirs.TryAdd(parts[1], 0)) + { + resumedDirs++; + } + break; + case "C": + if (_completedComputers.TryAdd(parts[1], 0)) + { + resumedComputers++; + } + break; + case "S": + if (parts.Length == 3) + { + ConcurrentBag bag = _computerShares.GetOrAdd(parts[1], _ => new ConcurrentBag()); + bag.Add(parts[2]); + resumedShares++; + } + break; + } + } + } + + mq.Info(string.Format( + "Resuming from checkpoint '{0}': {1} directories, {2} computers, and {3} shares already accounted for.", + _checkpointPath, resumedDirs, resumedComputers, resumedShares)); + } + else + { + mq.Info("Checkpoint file '" + _checkpointPath + "' doesn't exist yet, starting a fresh scan and recording progress to it."); + } + + // open for append; each write below is flushed immediately so progress survives a hard kill. + FileStream appendStream = new FileStream(_checkpointPath, FileMode.Append, FileAccess.Write, FileShare.Read); + _writer = new StreamWriter(appendStream, Encoding.UTF8) { AutoFlush = true }; + } + + public void Flush() + { + if (!Enabled || _writer == null) + { + return; + } + lock (_writeLock) + { + try + { + _writer.Flush(); + } + catch (Exception e) + { + BlockingMq.GetMq().Degub("Failed to flush checkpoint file: " + e.Message); + } + } + } + + public void Close() + { + if (!Enabled || _writer == null) + { + return; + } + lock (_writeLock) + { + try + { + _writer.Flush(); + _writer.Close(); + _writer = null; + } + catch (Exception e) + { + BlockingMq.GetMq().Degub("Failed to close checkpoint file: " + e.Message); + } + } + } + + private void WriteLine(string line) + { + if (!Enabled || _writer == null) + { + return; + } + lock (_writeLock) + { + try + { + _writer.WriteLine(line); + } + catch (Exception e) + { + // don't let checkpoint I/O problems take down the scan itself + BlockingMq.GetMq().Degub("Failed to write to checkpoint file: " + e.Message); + } + } + } + + public bool IsDirComplete(string dir) + { + return Enabled && _completedDirs.ContainsKey(dir); + } + + public void MarkDirComplete(string dir) + { + if (!Enabled) + { + return; + } + if (_completedDirs.TryAdd(dir, 0)) + { + WriteLine("D|" + dir); + } + } + + /// + /// Call once per directory, right after Directory.GetFiles() succeeds, with the + /// number of files found there. A directory is only checkpointed as complete once + /// every one of those files has actually finished being scanned (see + /// FileScanFinished) - NOT when the scan tasks are merely queued. Tree-walking + /// (listing names) races far ahead of content-scanning (reading/grepping bytes), + /// so marking completion at queue-time would checkpoint directories whose files + /// were never actually scanned, and a resumed run would skip them entirely. + /// + public void BeginDirFiles(string dir, int fileCount) + { + if (!Enabled) + { + return; + } + if (fileCount <= 0) + { + // no files in this dir to wait on, safe to checkpoint immediately. + MarkDirComplete(dir); + return; + } + _pendingFileCounts[dir] = fileCount; + } + + /// + /// Call exactly once per file that was counted in BeginDirFiles, when that file's + /// scan attempt is done (success OR failure - a permanently-failing file shouldn't + /// block the directory from ever checkpointing). Once every file counted for a + /// directory has reported in, the directory is marked complete. + /// + public void FileScanFinished(string dir) + { + if (!Enabled) + { + return; + } + int remaining = _pendingFileCounts.AddOrUpdate(dir, 0, (_, current) => current - 1); + if (remaining <= 0) + { + int dummy; + _pendingFileCounts.TryRemove(dir, out dummy); + MarkDirComplete(dir); + } + } + + public bool IsComputerComplete(string computer) + { + return Enabled && _completedComputers.ContainsKey(computer); + } + + public void MarkComputerComplete(string computer) + { + if (!Enabled) + { + return; + } + if (_completedComputers.TryAdd(computer, 0)) + { + WriteLine("C|" + computer); + } + } + + public void RecordShare(string computer, string sharePath) + { + if (!Enabled) + { + return; + } + ConcurrentBag bag = _computerShares.GetOrAdd(computer, _ => new ConcurrentBag()); + bag.Add(sharePath); + WriteLine("S|" + computer + "|" + sharePath); + } + + public IEnumerable GetComputerShares(string computer) + { + ConcurrentBag bag; + if (_computerShares.TryGetValue(computer, out bag)) + { + return bag; + } + return new string[0]; + } + } +} diff --git a/SnaffCore/Config/Options.cs b/SnaffCore/Config/Options.cs index 78c66282..8d6cc80d 100644 --- a/SnaffCore/Config/Options.cs +++ b/SnaffCore/Config/Options.cs @@ -32,6 +32,11 @@ public partial class Options public string CurrentUser { get; set; } = WindowsIdentity.GetCurrent().Name; public string RuleDir { get; set; } + // Checkpoint/resume options. Point this at a file and Snaffler will + // record progress there and skip already-completed work on a re-run + // against the same file (e.g. after a dropped connection or a sleep). + public string CheckpointFile { get; set; } + public int TimeOut { get; set; } = 5; // Concurrency Options diff --git a/SnaffCore/ShareFind/ShareFinder.cs b/SnaffCore/ShareFind/ShareFinder.cs index 4d1395c1..f88ea7c9 100644 --- a/SnaffCore/ShareFind/ShareFinder.cs +++ b/SnaffCore/ShareFind/ShareFinder.cs @@ -30,6 +30,31 @@ public ShareFinder() internal void GetComputerShares(string computer) { + // Resume support: if we already fully enumerated this computer's shares on + // a previous (interrupted) run, don't hit the network again - just replay + // the shares we already found scannable and let TreeWalker's own resume + // logic (ScanState.IsDirComplete) figure out what's left to do inside them. + if (SnaffCore.Checkpoint.ScanState.Instance.IsComputerComplete(computer)) + { + Mq.Trace("Skipping share enumeration on " + computer + ", already completed in a previous run."); + foreach (string knownShare in SnaffCore.Checkpoint.ScanState.Instance.GetComputerShares(computer)) + { + TreeTaskScheduler.New(() => + { + try + { + TreeWalker.WalkTree(knownShare); + } + catch (Exception e) + { + Mq.Error("Exception in TreeWalker task for share " + knownShare); + Mq.Error(e.ToString()); + } + }); + } + return; + } + // find the shares HostShareInfo[] hostShareInfos = GetHostShareInfo(computer); @@ -183,6 +208,10 @@ internal void GetComputerShares(string computer) if (MyOptions.ScanFoundShares) { + // remember this share so a resumed run can re-queue it without + // re-hitting the network to enumerate shares on this computer. + SnaffCore.Checkpoint.ScanState.Instance.RecordShare(computer, shareResult.SharePath); + Mq.Trace("Creating a TreeWalker task for " + shareResult.SharePath); TreeTaskScheduler.New(() => { @@ -206,6 +235,10 @@ internal void GetComputerShares(string computer) } } } + + // Share enumeration for this computer is done; a resumed run can skip + // straight to replaying the shares we already recorded above. + SnaffCore.Checkpoint.ScanState.Instance.MarkComputerComplete(computer); } internal bool IsShareReadable(string share) diff --git a/SnaffCore/SnaffCon.cs b/SnaffCore/SnaffCon.cs index 5334939d..05c30e68 100644 --- a/SnaffCore/SnaffCon.cs +++ b/SnaffCore/SnaffCon.cs @@ -42,6 +42,8 @@ public SnaffCon(Options options) MyOptions = options; Mq = BlockingMq.GetMq(); + SnaffCore.Checkpoint.ScanState.Instance.Load(MyOptions.CheckpointFile); + int shareThreads = MyOptions.ShareThreads; int treeThreads = MyOptions.TreeThreads; int fileThreads = MyOptions.FileThreads; @@ -156,6 +158,7 @@ public void Execute() TimeSpan runSpan = finished.Subtract(StartTime); Mq.Info("Finished at " + finished.ToLocalTime()); Mq.Info("Snafflin' took " + runSpan); + SnaffCore.Checkpoint.ScanState.Instance.Close(); Mq.Finish(); } diff --git a/SnaffCore/SnaffCore.csproj b/SnaffCore/SnaffCore.csproj index fa0fc804..ba598482 100644 --- a/SnaffCore/SnaffCore.csproj +++ b/SnaffCore/SnaffCore.csproj @@ -9,7 +9,7 @@ Properties SnaffCore SnaffCore - v4.5.1 + v4.8 512 true @@ -81,6 +81,7 @@ + diff --git a/SnaffCore/TreeWalk/TreeWalker.cs b/SnaffCore/TreeWalk/TreeWalker.cs index f4bf9c2b..fce34932 100644 --- a/SnaffCore/TreeWalk/TreeWalker.cs +++ b/SnaffCore/TreeWalk/TreeWalker.cs @@ -60,46 +60,66 @@ public void WalkTree(string currentDir) } - // Existing code: - try + // Resume support: if this directory's files were already fully queued for + // scanning on a previous (interrupted) run, don't do it again - just fall + // through to the (cheap, no file I/O) subdirectory recursion below so we + // keep walking down to wherever the previous run actually left off. + bool alreadyDone = SnaffCore.Checkpoint.ScanState.Instance.IsDirComplete(currentDir); + + if (!alreadyDone) { - string[] files = Directory.GetFiles(currentDir); - // check if we actually like the files - foreach (string file in files) + // Existing code: + try { - FileTaskScheduler.New(() => + string[] files = Directory.GetFiles(currentDir); + + // Register how many files this directory needs scanned before it can be + // checkpointed as complete. If there are zero, this marks it complete + // immediately; otherwise completion is reported per-file below as each + // FileScanner task actually finishes (not merely gets queued). + SnaffCore.Checkpoint.ScanState.Instance.BeginDirFiles(currentDir, files.Length); + + // check if we actually like the files + foreach (string file in files) { - try - { - FileScanner.ScanFile(file); - } - catch (Exception e) + FileTaskScheduler.New(() => { - Mq.Error("Exception in FileScanner task for file " + file); - Mq.Trace(e.ToString()); - } - }); + try + { + FileScanner.ScanFile(file); + } + catch (Exception e) + { + Mq.Error("Exception in FileScanner task for file " + file); + Mq.Trace(e.ToString()); + } + finally + { + SnaffCore.Checkpoint.ScanState.Instance.FileScanFinished(currentDir); + } + }); + } + } + catch (UnauthorizedAccessException) + { + //Mq.Trace(e.ToString()); + //continue; + } + catch (DirectoryNotFoundException) + { + //Mq.Trace(e.ToString()); + //continue; + } + catch (IOException) + { + //Mq.Trace(e.ToString()); + //continue; + } + catch (Exception e) + { + Mq.Degub(e.ToString()); + //continue; } - } - catch (UnauthorizedAccessException) - { - //Mq.Trace(e.ToString()); - //continue; - } - catch (DirectoryNotFoundException) - { - //Mq.Trace(e.ToString()); - //continue; - } - catch (IOException) - { - //Mq.Trace(e.ToString()); - //continue; - } - catch (Exception e) - { - Mq.Degub(e.ToString()); - //continue; } try @@ -182,14 +202,27 @@ public void WalkSccmTree(string currentDir, string sccmBaseDir) return; } + bool sccmDirAlreadyDone = SnaffCore.Checkpoint.ScanState.Instance.IsDirComplete(currentDir); + + if (!sccmDirAlreadyDone) + { try { string[] files = Directory.GetFiles(currentDir); + + SnaffCore.Checkpoint.ScanState.Instance.BeginDirFiles(currentDir, files.Length); + // check if we actually like the files foreach (string file in files) { FileTaskScheduler.New(() => { + // Scanning for this file is considered "done" either when this + // outer task bails out early, or - if it hands off to the inner + // FileScanner task below - when that inner task finishes instead. + // handedOff tracks which of those applies so exactly one of them + // reports completion back to ScanState. + bool handedOff = false; try { //FileScanner.ScanFile(file); @@ -245,6 +278,7 @@ public void WalkSccmTree(string currentDir, string sccmBaseDir) } */ + handedOff = true; FileTaskScheduler.New(() => { try @@ -256,6 +290,10 @@ public void WalkSccmTree(string currentDir, string sccmBaseDir) Mq.Error("Exception in FileScanner task for file " + file); Mq.Trace(e.ToString()); } + finally + { + SnaffCore.Checkpoint.ScanState.Instance.FileScanFinished(currentDir); + } }); @@ -286,6 +324,13 @@ public void WalkSccmTree(string currentDir, string sccmBaseDir) { Mq.Trace(e.ToString()); } + finally + { + if (!handedOff) + { + SnaffCore.Checkpoint.ScanState.Instance.FileScanFinished(currentDir); + } + } }); } } @@ -309,6 +354,7 @@ public void WalkSccmTree(string currentDir, string sccmBaseDir) Mq.Degub(e.ToString()); //continue; } + } // end !sccmDirAlreadyDone try { diff --git a/SnaffCore/UltraSnaffCore.csproj b/SnaffCore/UltraSnaffCore.csproj index 539d3311..e998623c 100644 --- a/SnaffCore/UltraSnaffCore.csproj +++ b/SnaffCore/UltraSnaffCore.csproj @@ -81,6 +81,7 @@ + diff --git a/Snaffler/Config.cs b/Snaffler/Config.cs index e9aedb21..531e1881 100644 --- a/Snaffler/Config.cs +++ b/Snaffler/Config.cs @@ -106,7 +106,9 @@ 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"); - // list of letters i haven't used yet: gnqw + ValueArgument checkpointArg = new ValueArgument('w', "checkpoint", + "Path to a checkpoint file used to resume an interrupted scan. First run against a given path creates it and records progress as it goes; if the scan gets cut off (network drop, sleep, crash, Ctrl-C) just run the exact same command again and Snaffler will skip everything already completed and carry on from there."); + // list of letters i haven't used yet: gnq CommandLineParser.CommandLineParser parser = new CommandLineParser.CommandLineParser(); parser.Arguments.Add(timeOutArg); @@ -131,7 +133,8 @@ private static Options ParseImpl(string[] args) parser.Arguments.Add(compTargetArg); parser.Arguments.Add(ruleDirArg); parser.Arguments.Add(logType); - parser.Arguments.Add(compExclusionArg); + parser.Arguments.Add(compExclusionArg); + parser.Arguments.Add(checkpointArg); // 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)) @@ -344,6 +347,12 @@ private static Options ParseImpl(string[] args) } } + if (checkpointArg.Parsed && !String.IsNullOrWhiteSpace(checkpointArg.Value)) + { + parsedConfig.CheckpointFile = checkpointArg.Value; + Mq.Info("Using checkpoint file: " + parsedConfig.CheckpointFile); + } + if (maxGrepSizeArg.Parsed) { parsedConfig.MaxSizeToGrep = maxGrepSizeArg.Value; @@ -400,9 +409,23 @@ private static Options ParseImpl(string[] args) string configFile = configFileArg.Value; parsedConfig = Toml.ReadFile(configFile, settings); Mq.Info("Read config file from " + configFile); + + // Loading a toml config replaces parsedConfig wholesale, which would + // silently drop a -w/--checkpoint value parsed earlier. Re-apply it so + // the CLI flag always wins. + if (checkpointArg.Parsed && !String.IsNullOrWhiteSpace(checkpointArg.Value)) + { + parsedConfig.CheckpointFile = checkpointArg.Value; + } } } + if (!String.IsNullOrWhiteSpace(parsedConfig.CheckpointFile)) + { + string resolvedCheckpointPath = System.IO.Path.GetFullPath(parsedConfig.CheckpointFile); + Console.WriteLine("[Checkpoint] Using checkpoint file: " + resolvedCheckpointPath); + } + if (!parsedConfig.LogToConsole && !parsedConfig.LogToFile) { Mq.Error( diff --git a/Snaffler/Properties/Resources.Designer.cs b/Snaffler/Properties/Resources.Designer.cs index 01383f00..152c05a6 100644 --- a/Snaffler/Properties/Resources.Designer.cs +++ b/Snaffler/Properties/Resources.Designer.cs @@ -19,7 +19,7 @@ namespace Snaffler.Properties { // class via a tool like ResGen or Visual Studio. // To add or remove a member, edit your .ResX file then rerun ResGen // with the /str option, or rebuild your VS project. - [global::System.CodeDom.Compiler.GeneratedCodeAttribute("System.Resources.Tools.StronglyTypedResourceBuilder", "16.0.0.0")] + [global::System.CodeDom.Compiler.GeneratedCodeAttribute("System.Resources.Tools.StronglyTypedResourceBuilder", "17.0.0.0")] [global::System.Diagnostics.DebuggerNonUserCodeAttribute()] [global::System.Runtime.CompilerServices.CompilerGeneratedAttribute()] public class Resources { diff --git a/Snaffler/SnaffleRunner.cs b/Snaffler/SnaffleRunner.cs index e5f32497..b9e4fcb3 100644 --- a/Snaffler/SnaffleRunner.cs +++ b/Snaffler/SnaffleRunner.cs @@ -44,6 +44,20 @@ public void Run(string[] args) hostString(); // print the thing PrintBanner(); + + // Make sure Ctrl-C doesn't hard-kill the process before the checkpoint + // writer gets a chance to flush. The checkpoint log is written immediately + // as each unit of work completes (AutoFlush), so this is mostly a safety + // net, but it also gives a clear on-screen confirmation of where things + // stand when you interrupt a scan. + Console.CancelKeyPress += (sender, eventArgs) => + { + Console.WriteLine("\n[Checkpoint] Ctrl-C caught, flushing checkpoint state before exiting..."); + SnaffCore.Checkpoint.ScanState.Instance.Close(); + Console.WriteLine("[Checkpoint] Done. Re-run the same command (same -w path) to resume."); + eventArgs.Cancel = false; // let the process actually terminate now + }; + // set up the message queue for operation BlockingMq.MakeMq(); // get a handle to the message queue singleton diff --git a/Snaffler/Snaffler.csproj b/Snaffler/Snaffler.csproj index f84fdba9..787ee386 100644 --- a/Snaffler/Snaffler.csproj +++ b/Snaffler/Snaffler.csproj @@ -8,7 +8,7 @@ Exe Snaffler Snaffler - v4.5.1 + v4.8 512 true diff --git a/Snaffler/app.config b/Snaffler/app.config index 51278a45..3e0e37cf 100644 --- a/Snaffler/app.config +++ b/Snaffler/app.config @@ -1,3 +1,3 @@ - + diff --git a/snafflertest/german-config/Kennw#U00f6rter.txt b/snafflertest/german-config/Kennw#U00f6rter.txt new file mode 100644 index 00000000..e69de29b diff --git a/snafflertest/german-config/Schl#U00fcssel.txt b/snafflertest/german-config/Schl#U00fcssel.txt new file mode 100644 index 00000000..e69de29b