From 8defc94131c082a42cdfc82e773739af617a9458 Mon Sep 17 00:00:00 2001
From: Rohan Vazarkar <rvazarkr@gmail.com>
Date: Thu, 3 Nov 2022 11:15:56 -0400
Subject: [PATCH] fix: add timeout to registry enum
 (https://github.com/BloodHoundAD/BloodHound/issues/541) fix: lookup accounts
 in registry enum to remove local ones

Closes: https://github.com/BloodHoundAD/SharpHoundCommon/issues/27
---
 .../Processors/ComputerSessionProcessor.cs    | 42 +++++++++++++++----
 1 file changed, 35 insertions(+), 7 deletions(-)

diff --git a/src/CommonLib/Processors/ComputerSessionProcessor.cs b/src/CommonLib/Processors/ComputerSessionProcessor.cs
index ed96685b..d4760703 100644
--- a/src/CommonLib/Processors/ComputerSessionProcessor.cs
+++ b/src/CommonLib/Processors/ComputerSessionProcessor.cs
@@ -228,7 +228,7 @@ public SessionAPIResult ReadUserSessionsPrivileged(string computerName,
             return ret;
         }
 
-        public SessionAPIResult ReadUserSessionsRegistry(string computerName, string computerDomain,
+        public async Task<SessionAPIResult> ReadUserSessionsRegistry(string computerName, string computerDomain,
             string computerSid)
         {
             var ret = new SessionAPIResult();
@@ -237,7 +237,24 @@ public SessionAPIResult ReadUserSessionsRegistry(string computerName, string com
 
             try
             {
-                key = RegistryKey.OpenRemoteBaseKey(RegistryHive.Users, computerName);
+                var task = OpenRegistryKey(computerName, RegistryHive.Users);
+                
+                if (await Task.WhenAny(task, Task.Delay(10000)) != task)
+                {
+                    _log.LogDebug("Hit timeout on registry enum on {Server}. Abandoning registry enum", computerName);
+                    ret.Collected = false;
+                    ret.FailureReason = "Timeout";
+                    SendComputerStatus(new CSVComputerStatus
+                    {
+                        Status = "Timeout",
+                        Task = "RegistrySessionEnum",
+                        ComputerName = computerName
+                    });
+                    return ret;
+                }
+                
+                key = task.Result;
+
                 ret.Collected = true;
                 SendComputerStatus(new CSVComputerStatus
                 {
@@ -246,11 +263,17 @@ public SessionAPIResult ReadUserSessionsRegistry(string computerName, string com
                     ComputerName = computerName
                 });
                 _log.LogDebug("Registry session enum succeeded on {ComputerName}", computerName);
-                ret.Results = key.GetSubKeyNames().Where(subkey => SidRegex.IsMatch(subkey)).Select(x => new Session
-                {
-                    ComputerSID = computerSid,
-                    UserSID = x
-                }).ToArray();
+                ret.Results = key.GetSubKeyNames()
+                    .Where(subkey => SidRegex.IsMatch(subkey))
+                    .Select(x => _utils.ResolveIDAndType(x, computerDomain))
+                    .Where(x => x != null)
+                    .Select(x =>
+                        new Session
+                        {
+                            ComputerSID = computerSid,
+                            UserSID = x.ObjectIdentifier
+                        })
+                    .ToArray();
 
                 return ret;
             }
@@ -273,6 +296,11 @@ public SessionAPIResult ReadUserSessionsRegistry(string computerName, string com
             }
         }
 
+        private Task<RegistryKey> OpenRegistryKey(string computerName, RegistryHive hive)
+        {
+            return Task.Run(() => RegistryKey.OpenRemoteBaseKey(hive, computerName));
+        }
+
         private void SendComputerStatus(CSVComputerStatus status)
         {
             ComputerStatusEvent?.Invoke(status);