Files
Flyoobe/FlybyScript/USBManagement.cs
T
2024-10-26 09:25:41 +02:00

441 lines
18 KiB
C#

using System;
using System.Diagnostics;
using System.Drawing;
using System.IO;
using System.Net.Http;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using System.Windows.Forms;
namespace FlybyScript
{
public class USBManagement
{
private readonly Logger _logger;
private ProgressBar _progressBar;
public Action<string> UpdateLabelStatus;
private int currentStep = 0;
private static readonly HttpClient httpClient = new HttpClient();
public USBManagement(Logger logger, Action<string> updateLabelStatus, ProgressBar progressBar)
{
_logger = logger;
_progressBar = progressBar;
UpdateLabelStatus = updateLabelStatus; // Set the delegate to update UI
}
private string[] steps = new[]
{
"Format and Boot USB Drive",
"Mount ISO",
"Copy ISO to USB",
"Split Install.wim",
"Finalize Installation Media (Install.wim)"
};
/// <summary>
/// Updates the LabelStatus with the current step
/// </summary>
public void UpdateStatusLabel()
{
currentStep++;
if (currentStep <= steps.Length)
{
string statusText = $"{currentStep}/{steps.Length} - {steps[currentStep - 1]}";
// Call the delegate to update the label in the user control
UpdateLabelStatus?.Invoke(statusText);
// Log to the logger as well
_logger.Log($"Step {currentStep}: {steps[currentStep - 1]} started.", Color.Blue);
}
}
/// <summary>
/// Updates the progress bar with the specified value
/// </summary>
private void UpdateProgress(int value)
{
if (_progressBar.InvokeRequired)
{
_progressBar.Invoke(new Action(() => _progressBar.Value = value));
}
else
{
_progressBar.Value = value;
}
}
/// <summary>
/// Formats the USB drive using DiskPart
/// </summary>
/// <param name="diskNumber">Diskpart Device ID</param>
/// <returns></returns>
public async Task FormatDriveAsync(string diskNumber)
{
await Task.Run(() =>
{
try
{
ProcessStartInfo processInfo = new ProcessStartInfo
{
FileName = "diskpart.exe",
RedirectStandardInput = true,
RedirectStandardOutput = true,
UseShellExecute = false,
CreateNoWindow = true
};
using (Process process = Process.Start(processInfo))
{
using (StreamWriter writer = process.StandardInput)
{
if (writer.BaseStream.CanWrite)
{
// Step 1: Format and Boot USB Drive
UpdateStatusLabel(); // Updates to 1/5 - Format and Boot USB Drive
// Send commands to DiskPart
writer.WriteLine($"select disk {diskNumber}"); // Select the disk based on user input
UpdateProgress(10); // Update progress
writer.WriteLine("clean"); // Clean the disk
writer.WriteLine("convert mbr"); // Convert to MBR
UpdateProgress(20);
writer.WriteLine("create partition primary"); // Create a new partition
writer.WriteLine("select partition 1"); // Select the new partition
writer.WriteLine("active"); // Set the new partition active
UpdateProgress(60);
writer.WriteLine("format fs=fat32 quick label=\"FlybyWin11\""); // Quick format to FAT32
// writer.WriteLine("assign letter=W"); // Assign drive letter
writer.WriteLine("exit"); // Exit DiskPart
UpdateProgress(100); // Final progress
}
}
string output = process.StandardOutput.ReadToEnd(); // Get the output from DiskPart
process.WaitForExit(); // Wait for the process to exit
_logger.Log($"Formatting completed: {output}", Color.Black);
}
}
catch (Exception ex)
{
_logger.Log($"Error formatting drive: {ex.Message}", Color.Red);
}
});
}
/// <summary>
/// Copies the contents of the specified ISO file to the USB drive using Robocopy.
/// </summary>
/// <param name="usbDrive">The drive letter of the USB drive.</param>
/// <param name="isoPath">The path to the ISO file.</param>
public async Task CopyISOAsync(string usbDrive, string isoPath)
{
try
{
// Step 1: Mount the ISO using PowerShell
UpdateStatusLabel(); // Updates to "2/5 - Mount ISO"
string mountCommand = $"Mount-DiskImage -ImagePath \"{isoPath}\"";
ProcessStartInfo mountInfo = new ProcessStartInfo("powershell", mountCommand)
{
RedirectStandardOutput = true,
UseShellExecute = false,
CreateNoWindow = true
};
using (Process mountProcess = Process.Start(mountInfo))
{
string output = await mountProcess.StandardOutput.ReadToEndAsync();
await mountProcess.WaitForExitAsync(); // Wait for the process to complete
_logger.Log($"ISO mounted: {output}", Color.Black);
}
// Step 2: Get the drive letter of the mounted ISO
string isoDriveLetter = await GetMountedDriveLetterAsync(); // Get the drive letter of the mounted ISO
if (string.IsNullOrEmpty(isoDriveLetter))
{
_logger.Log("Failed to get the mounted ISO drive letter.", Color.Red);
return; // Exit if we cant find the drive letter
}
_logger.Log($"Mounted ISO drive letter: {isoDriveLetter}", Color.Black);
_logger.Log($"USB drive letter: {usbDrive}", Color.Black);
// Step 3: Prepare the Robocopy command to copy everything except install.wim
string robocopyCommand = $"robocopy {isoDriveLetter.TrimEnd('\\')} {usbDrive.TrimEnd('\\')} /MIR /XF \"install.wim\"";
// Step 4: Execute the Robocopy command
ProcessStartInfo robocopyInfo = new ProcessStartInfo("cmd.exe", $"/C {robocopyCommand}")
{
RedirectStandardOutput = true,
UseShellExecute = false,
CreateNoWindow = true,
};
// Step 3: Copy ISO to USB
UpdateStatusLabel(); // Updates to "3/5 - Copy ISO to USB"
using (Process robocopyProcess = Process.Start(robocopyInfo))
{
robocopyProcess.OutputDataReceived += (s, e) =>
{
if (!string.IsNullOrEmpty(e.Data))
{
_logger.Log(e.Data, Color.Black);
// Calculate and update progress based on Robocopy output
int progress = CalculateProgressFromRobocopy(e.Data);
UpdateProgress(progress);
}
};
robocopyProcess.BeginOutputReadLine();
// Wait for the process to finish
await robocopyProcess.WaitForExitAsync();
}
// Step 5: Get the path to the original WIM file (from the ISO)
string wimFilePath = Path.Combine(isoDriveLetter, "sources", "install.wim");
_logger.Log("Preparing install.wim... This may take a moment, please wait.", Color.Blue);
// Step 4: Split Install.wim
UpdateStatusLabel(); // Updates to "4/5 - Split Install.wim"
if (File.Exists(wimFilePath))
{
// Step 7: Define the path for the split files
string swmFilePath = Path.Combine(usbDrive, "sources", "install.swm"); // First split file
// Step 8: Dism command to split WIM
string splitDismCommand = $"/C Dism /Split-Image /ImageFile:\"{wimFilePath}\" /SWMFile:\"{swmFilePath}\" /FileSize:3000";
// Execute the command
ProcessStartInfo splitInfo = new ProcessStartInfo("cmd.exe", splitDismCommand)
{
RedirectStandardOutput = true,
UseShellExecute = false,
CreateNoWindow = true
};
// Step 5: Finalize Installation Media
UpdateStatusLabel(); // Updates to 5/5 - Finalize Installation Media
using (Process splitProcess = Process.Start(splitInfo))
{
splitProcess.OutputDataReceived += (s, e) =>
{
if (!string.IsNullOrEmpty(e.Data))
{
_logger.Log(e.Data, Color.Black); // Log the output from Dism
}
};
splitProcess.BeginOutputReadLine();
// Wait for the process to finish
await splitProcess.WaitForExitAsync();
_logger.Log($"WIM file split completed with exit code: {splitProcess.ExitCode}", Color.Blue);
}
}
else
{
_logger.Log($"WIM file not found at {wimFilePath}.", Color.Red);
}
}
catch (Exception ex)
{
_logger.Log($"Error copying ISO: {ex.Message}", Color.Red);
}
MessageBox.Show("Your installation media is ready! You can now proceed with the installation.", "Installation Media", MessageBoxButtons.OK, MessageBoxIcon.Information);
}
/// <summary>
/// Retrieves the drive letter of the currently mounted drive.
/// </summary>
/// <returns>The drive letter of the mounted CD/DVD drive, or null if not found.</returns>
private async Task<string> GetMountedDriveLetterAsync()
{
try
{
// List all the drives to find the mounted CD/DVD drive
var drives = DriveInfo.GetDrives();
foreach (var drive in drives)
{
// Check if the drive is a CD/DVD and if it's mounted
if (drive.DriveType == DriveType.CDRom && drive.IsReady)
{
// Return the letter of the mounted drive
return drive.Name; // Returns "D:\", "E:\", etc.
}
}
return null; // Return null if no mounted CD/DVD drive is found
}
catch (Exception ex)
{
_logger.Log($"Error getting mounted drive letter: {ex.Message}", Color.Crimson);
return null;
}
}
/// <summary>
/// Downloads the Microsoft Media Creation Tool from https://go.microsoft.com/fwlink/?linkid=2156295
/// </summary>
/// <param name="downloadUrl">The URL of the Media Creation Tool.</param>
/// <returns>A task representing the asynchronous operation.</returns>
public async Task DownloadMediaCreationTool()
{
string downloadUrl = "https://go.microsoft.com/fwlink/?linkid=2156295";
using (SaveFileDialog saveFileDialog = new SaveFileDialog())
{
saveFileDialog.Filter = "Executable Files (*.exe)|*.exe";
saveFileDialog.Title = "Save Media Creation Tool";
saveFileDialog.FileName = "MediaCreationTool.exe";
if (saveFileDialog.ShowDialog() == DialogResult.OK)
{
string destinationPath = saveFileDialog.FileName;
try
{
_logger.Log("Preparing to download Media Creation Tool...", Color.Black);
_logger.Log("Download may take some time, please wait...", Color.Orange);
var response = await httpClient.GetAsync(downloadUrl);
response.EnsureSuccessStatusCode();
using (var fileStream = new FileStream(destinationPath, FileMode.Create, FileAccess.Write, FileShare.None))
{
await response.Content.CopyToAsync(fileStream);
}
_logger.Log("Download completed successfully!", Color.Green);
// Start the downloaded Media Creation Tool
_logger.Log("Starting Media Creation Tool...", Color.Black);
Process.Start(new ProcessStartInfo
{
FileName = destinationPath,
UseShellExecute = true // run .exe files
});
}
catch (Exception ex)
{
_logger.Log($"An error occurred while downloading: {ex.Message}", Color.Red);
}
}
else
{
_logger.Log("Download canceled by user.", Color.Red);
}
}
}
/// <summary>
/// Calculates progress based on Robocopy output
/// </summary>
private int CalculateProgressFromRobocopy(string robocopyOutput)
{
// Find the percentage completion in Robocopy output
var match = Regex.Match(robocopyOutput, @"(\d+)\%");
if (match.Success)
{
// Extract the percentage number from the match and convert it to an integer
if (int.TryParse(match.Groups[1].Value, out int progress))
{
return progress;
}
}
return 0; // Return 0 if no percentage was found
}
/// <summary>
/// Copy the ISO to the USB drive and optionally create an unattend.xml file
/// </summary>
/// <param name="usbDrive"></param>
/// <param name="isoPath"></param>
/// <param name="createUnattend"></param>
/// <returns></returns>
public async Task CopyISOWithUnattendAsync(string usbDrive, string isoPath, bool createUnattend)
{
await CopyISOAsync(usbDrive, isoPath);
if (createUnattend)
{
CreateUnattendXml(usbDrive); // Only create unattend.xml if checkbox is checked
}
}
/// <summary>
/// // Create unattend.xml file in the sources\$OEM$\$$\Panther directory
/// </summary>
/// <param name="usbDrive"></param>
public void CreateUnattendXml(string usbDrive)
{
string unattendDir = Path.Combine(usbDrive, "sources", "$OEM$", "$$", "Panther");
Directory.CreateDirectory(unattendDir); // Create the directory if it doesn't exist
string unattendPath = Path.Combine(unattendDir, "unattend.xml");
// Check if the file already exists
if (!File.Exists(unattendPath))
{
// Create the unattend.xml content
string xmlContent = @"<unattend xmlns='urn:schemas-microsoft-com:unattend'>
<settings pass='disabled'>
<component xmlns:wcm='http://schemas.microsoft.com/WMIConfig/2002/State' xmlns:xsi='http://www.w3.org/2001/XMLSchema-instance' name='Microsoft-Windows-Setup' processorArchitecture='amd64' language='neutral' publicKeyToken='31bf3856ad364e35' versionScope='nonSxS'>
<UserData>
<ProductKey>
<Key/>
</ProductKey>
</UserData>
<RunSynchronous>
<RunSynchronousCommand wcm:action='add'>
<Order>1</Order>
<Path>reg add HKLM\SYSTEM\Setup\LabConfig /v BypassTPMCheck /t REG_DWORD /d 1 /f</Path>
</RunSynchronousCommand>
<RunSynchronousCommand wcm:action='add'>
<Order>2</Order>
<Path>reg add HKLM\SYSTEM\Setup\LabConfig /v BypassSecureBootCheck /t REG_DWORD /d 1 /f</Path>
</RunSynchronousCommand>
<RunSynchronousCommand wcm:action='add'>
<Order>3</Order>
<Path>reg add HKLM\SYSTEM\Setup\LabConfig /v BypassRAMCheck /t REG_DWORD /d 1 /f</Path>
</RunSynchronousCommand>
</RunSynchronous>
</component>
</settings>
<settings pass='specialize'>
<component xmlns:wcm='http://schemas.microsoft.com/WMIConfig/2002/State' xmlns:xsi='http://www.w3.org/2001/XMLSchema-instance' name='Microsoft-Windows-Deployment' processorArchitecture='amd64' language='neutral' publicKeyToken='31bf3856ad364e35' versionScope='nonSxS'>
<RunSynchronous>
<RunSynchronousCommand wcm:action='add'>
<Order>1</Order>
<Path>reg add HKLM\SOFTWARE\Microsoft\Windows\CurrentVersion\OOBE /v BypassNRO /t REG_DWORD /d 1 /f</Path>
</RunSynchronousCommand>
</RunSynchronous>
</component>
</settings>
</unattend>";
File.WriteAllText(unattendPath, xmlContent);
_logger.Log("Created unattend.xml in the sources\\$OEM$\\$$\\Panther directory.", Color.Green);
}
else
{
_logger.Log("unattend.xml already exists in the sources\\$OEM$\\$$\\Panther directory.", Color.Blue);
}
}
}
}