wip: remove some dead code, add some missing functions

SpecterOps · Jul 2, 2024 · 4673bbd · 4673bbd
1 parent ef2b38e
commit 4673bbd
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 158 deletions.
diff --git a/src/CommonLib/DCConnectionCache.cs b/src/CommonLib/DCConnectionCache.cs
diff --git a/src/CommonLib/ILdapUtilsNew.cs b/src/CommonLib/ILdapUtilsNew.cs
@@ -42,5 +42,7 @@ IAsyncEnumerable<Result<string>> RangedRetrieval(string distinguishedName,
 
         public Task<bool> IsDomainController(string computerObjectId, string domainName);
         public Task<(bool Success, TypedPrincipal Principal)> LookupDistinguishedName(string distinguishedName);
+        public void AddDomainController(string domainControllerSID);
+        IAsyncEnumerable<OutputBase> GetWellKnownPrincipalOutput();
     }
 }
diff --git a/src/CommonLib/LdapUtilsNew.cs b/src/CommonLib/LdapUtilsNew.cs
@@ -29,6 +29,7 @@ public class LdapUtilsNew : ILdapUtilsNew {
         //This cache is indexed by domain sid
         private readonly ConcurrentDictionary<string, NetAPIStructs.DomainControllerInfo?> _dcInfoCache = new();
         private static readonly ConcurrentDictionary<string, Domain> DomainCache = new();
+        private static readonly ConcurrentDictionary<string, byte> DomainControllers = new();
 
         private static readonly ConcurrentDictionary<string, string> DomainToForestCache =
             new(StringComparer.OrdinalIgnoreCase);
@@ -165,22 +166,31 @@ await _connectionPool.GetLdapConnection(domain,
                             _log.LogError(
                                 "RangedRetrieval - Failed to get a new connection after ServerDown for path {Path}",
                                 distinguishedName);
+                            _connectionPool.ReleaseConnection(connectionWrapper);
                             tempResult =
                                 Result<string>.Fail(
                                     "RangedRetrieval - Failed to get a new connection after ServerDown.");
                         }
                     }
                 }
                 catch (LdapException le) {
+                    if (le.ErrorCode is (int)LdapErrorCodes.ServerDown) {
+                        _connectionPool.ReleaseConnection(connectionWrapper, true);
+                    }
+                    else {
+                        _connectionPool.ReleaseConnection(connectionWrapper);
+                    }
                     tempResult = Result<string>.Fail(
                         $"Caught unrecoverable ldap exception: {le.Message} (ServerMessage: {le.ServerErrorMessage}) (ErrorCode: {le.ErrorCode})");
                 }
                 catch (Exception e) {
+                    _connectionPool.ReleaseConnection(connectionWrapper);
                     tempResult =
                         Result<string>.Fail($"Caught unrecoverable exception: {e.Message}");
                 }
 
                 //If we have a tempResult set it means we hit an error we couldn't recover from, so yield that result and then break out of the function
+                //We handle connection release in the relevant exception blocks
                 if (tempResult != null) {
                     yield return tempResult;
                     yield break;
@@ -201,6 +211,7 @@ await _connectionPool.GetLdapConnection(domain,
                     }
 
                     if (complete) {
+                        _connectionPool.ReleaseConnection(connectionWrapper);
                         yield break;
                     }
 
@@ -210,6 +221,7 @@ await _connectionPool.GetLdapConnection(domain,
                 }
                 else {
                     //I dont know what can cause a RR to have multiple entries, but its nothing good. Break out
+                    _connectionPool.ReleaseConnection(connectionWrapper);
                     yield break;
                 }
             }
@@ -1382,5 +1394,38 @@ public async Task<bool> IsDomainController(string computerObjectId, string domai
                 }
             }
         }
+
+        public void AddDomainController(string domainControllerSID)
+        {
+            DomainControllers.TryAdd(domainControllerSID, new byte());
+        }
+
+        public async IAsyncEnumerable<OutputBase> GetWellKnownPrincipalOutput() {
+            foreach (var wkp in SeenWellKnownPrincipals)
+            {
+                WellKnownPrincipal.GetWellKnownPrincipal(wkp.Value.WkpId, out var principal);
+                OutputBase output = principal.ObjectType switch
+                {
+                    Label.User => new User(),
+                    Label.Computer => new Computer(),
+                    Label.Group => new OutputTypes.Group(),
+                    Label.GPO => new GPO(),
+                    Label.Domain => new OutputTypes.Domain(),
+                    Label.OU => new OU(),
+                    Label.Container => new Container(),
+                    Label.Configuration => new Container(),
+                    _ => throw new ArgumentOutOfRangeException()
+                };
+
+                output.Properties.Add("name", $"{principal.ObjectIdentifier}@{wkp.Value.DomainName}".ToUpper());
+                if (await GetDomainSidFromDomainName(wkp.Value.DomainName) is (true, var sid)) {
+                    output.Properties.Add("domainsid", sid);    
+                }
+
+                output.Properties.Add("domain", wkp.Value.DomainName.ToUpper());
+                output.ObjectIdentifier = wkp.Key;
+                yield return output;
+            }
+        }
     }
 }
diff --git a/src/CommonLib/SearchResultEntryWrapperNew.cs b/src/CommonLib/SearchResultEntryWrapperNew.cs