fix: set collected correctly for domains

SpecterOps · Jan 18, 2024 · cf35034 · cf35034
1 parent 17f6f6d
commit cf35034
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 9 deletions.
diff --git a/src/BaseContext.cs b/src/BaseContext.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
@@ -64,6 +65,8 @@ public void UpdateLoopTime()
             CurrentLoopTime = $"{DateTime.Now:yyyyMMddHHmmss}";
         }
 
+        public HashSet<string> CollectedDomainSids { get; } = new();
+
         public async Task DoDelay()
         {
             if (Throttle == 0)

diff --git a/src/Client/Context.cs b/src/Client/Context.cs
@@ -75,5 +75,6 @@ public interface IContext
         string ResolveFileName(string filename, string extension, bool addTimestamp);
         EnumerationDomain[] Domains { get; set; }
         void UpdateLoopTime();
+        public HashSet<string> CollectedDomainSids { get; }
     }
 }
diff --git a/src/Producers/LdapProducer.cs b/src/Producers/LdapProducer.cs
@@ -44,7 +44,7 @@ public override async Task Produce()
 
             foreach (var domain in Context.Domains)
             {
-                Context.Logger.LogInformation("Beginning LDAP search for {Domain}", domain);
+                Context.Logger.LogInformation("Beginning LDAP search for {Domain}", domain.Name);
                 //Do a basic  LDAP search and grab results
                 var successfulConnect = false;
                 try
@@ -64,14 +64,7 @@ public override async Task Produce()
                     continue;
                 }
 
-                await OutputChannel.Writer.WriteAsync(new Domain
-                {
-                    ObjectIdentifier = domain.DomainSid,
-                    Properties = new Dictionary<string, object>
-                    {
-                        { "collected", true },
-                    }
-                });
+                Context.CollectedDomainSids.Add(domain.DomainSid);
 
                 foreach (var searchResult in Context.LDAPUtils.QueryLDAP(ldapData.Filter.GetFilter(), SearchScope.Subtree,
                              ldapData.Props.Distinct().ToArray(), cancellationToken, domain.Name,

diff --git a/src/Runtime/LDAPConsumer.cs b/src/Runtime/LDAPConsumer.cs
@@ -37,6 +37,11 @@ internal static async Task ConsumeSearchResults(Channel<ISearchResultEntry> inpu
                         watch.Elapsed.TotalMilliseconds, res.DisplayName);
                     if (processed == null)
                         continue;
+
+                    if (processed is Domain d && context.CollectedDomainSids.Contains(d.ObjectIdentifier))
+                    {
+                        d.Properties.Add("collected", true);
+                    }
                     await outputChannel.Writer.WriteAsync(processed);
                 }
                 catch (Exception e)