Write files in windowsPE pass more efficiently

#99
cschneegans · Nov 5, 2024 · 499f40f · 499f40f
1 parent 31f7829
commit 499f40f
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 29 deletions.
diff --git a/Extensions.cs b/Extensions.cs
@@ -46,4 +46,14 @@ public static IImmutableDictionary<string, T> ToKeyedDictionary<T>(this IEnumera
       keyComparer: StringComparer.OrdinalIgnoreCase
     );
   }
+
+  public static string JoinString<T>(this IEnumerable<T> enumerable, string separator)
+  {
+    return string.Join(separator, enumerable);
+  }
+
+  public static string JoinString<T>(this IEnumerable<T> enumerable, char separator)
+  {
+    return string.Join(separator, enumerable);
+  }
 }
diff --git a/Main.cs b/Main.cs
@@ -217,8 +217,43 @@ public static string InvokeJScript(string filepath)
     return @$"cscript.exe //E:jscript ""{filepath}""";
   }
 
-  public static IEnumerable<string> WriteToFile(string path, string line)
+  public static List<string> WriteToFile(string path, IEnumerable<string> lines)
   {
+    static IEnumerable<string> Trim(IEnumerable<string> input)
+    {
+      return input
+        .Select(l => l.Trim())
+        .Where(e => e.Length > 0);
+    }
+
+    static IEnumerable<IEnumerable<string>> Join(IEnumerable<string> input)
+    {
+      int max = 200;
+      List<List<string>> output = [];
+      List<string> current = [];
+      foreach (string line in input)
+      {
+        if (current.Count == 0)
+        {
+          current.Add(line);
+        }
+        else if (current.Sum(e => e.Length + 10) + line.Length > max)
+        {
+          output.Add(current);
+          current = [line];
+        }
+        else
+        {
+          current.Add(line);
+        }
+      }
+      if (current.Count > 0)
+      {
+        output.Add(current);
+      }
+      return output;
+    }
+
     static string EscapeShell(string command)
     {
       return command
@@ -227,20 +262,24 @@ static string EscapeShell(string command)
         .Replace("<", "^<")
         .Replace(">", "^>")
         .Replace("|", "^|")
-        .Replace("%", "^%");
+        .Replace("%", "^%")
+        .Replace(")", "^)");
     }
 
-    if (string.IsNullOrWhiteSpace(line))
+    IEnumerable<string> Enumerate()
     {
-      yield break;
-    }
-
-    string command = $@"cmd.exe /c "">>""{path}"" echo {EscapeShell(line)}""";
-    if (command.Length >= 255)
-    {
-      throw new ConfigurationException($"Line '{line}' is too long.");
+      foreach (IEnumerable<string> group in Join(Trim(lines)))
+      {
+        string echos = group.Select(l => $"echo {EscapeShell(l)}").JoinString('&');
+        string command = $@"cmd.exe /c "">>""{path}"" ({echos})""";
+        if (command.Length >= 255)
+        {
+          throw new ConfigurationException($"Line '{command}' is too long.");
+        }
+        yield return command;
+      }
     }
-    yield return command;
+    return Enumerate().ToList();
   }
 }
 

diff --git a/modifier/Disk.cs b/modifier/Disk.cs
@@ -132,15 +132,10 @@ private void WriteScript(IEnumerable<string> lines)
     string log = @"X:\diskpart.log";
 
     CommandAppender appender = GetAppender(CommandConfig.WindowsPE);
-    foreach (string line in lines)
-    {
-      appender.Append(
-        CommandBuilder.WriteToFile(script, line)
-      );
-    }
-    appender.Append(
-      CommandBuilder.ShellCommand($@"diskpart.exe /s ""{script}"" >>""{log}"" || ( type ""{log}"" & echo diskpart encountered an error. & pause & exit /b 1 )")
-    );
+    appender.Append([
+      ..CommandBuilder.WriteToFile(script, lines),
+      CommandBuilder.ShellCommand($@"diskpart.exe /s ""{script}"" >>""{log}"" || ( type ""{log}"" & echo diskpart encountered an error. & pause & exit /b 1 )"),
+    ]);
   }
 
   public static string GetCustomDiskpartScript()

diff --git a/modifier/Optimizations.cs b/modifier/Optimizations.cs
@@ -146,16 +146,11 @@ IEnumerable<string> SetExplorerOptions(string rootKey, string subKey)
     if (Configuration.DisableDefender)
     {
       CommandAppender pe = GetAppender(CommandConfig.WindowsPE);
-      const string path = @"X:\disable-defender.vbs";
-      foreach (string line in Util.SplitLines(Util.StringFromResource("disable-defender.vbs")))
-      {
-        pe.Append(
-          CommandBuilder.WriteToFile(path, line)
-        );
-      }
-      pe.Append(
+      const string path = @"X:\defender.vbs";
+      pe.Append([
+        ..CommandBuilder.WriteToFile(path, Util.SplitLines(Util.StringFromResource("disable-defender.vbs"))),
         CommandBuilder.ShellCommand($"start /MIN cscript.exe //E:vbscript {path}")
-      );
+      ]);
     }
 
     if (Configuration.DisableSac)