383 lines
13 KiB
Rust
383 lines
13 KiB
Rust
// This file is part of Deja-vu.
|
|
//
|
|
// Deja-vu is free software: you can redistribute it and/or modify it
|
|
// under the terms of the GNU General Public License as published by
|
|
// the Free Software Foundation, either version 3 of the License, or
|
|
// (at your option) any later version.
|
|
//
|
|
// Deja-vu is distributed in the hope that it will be useful, but
|
|
// WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
|
|
// See the GNU General Public License for more details.
|
|
//
|
|
// You should have received a copy of the GNU General Public License
|
|
// along with Deja-vu. If not, see <https://www.gnu.org/licenses/>.
|
|
//
|
|
use std::{
|
|
collections::HashMap,
|
|
fs::File,
|
|
io::Write,
|
|
process::{Command, Output, Stdio},
|
|
};
|
|
|
|
use once_cell::sync::Lazy;
|
|
use regex::Regex;
|
|
use tempfile::tempdir;
|
|
use tracing::{info, warn};
|
|
|
|
use crate::system::disk::{
|
|
get_disk_serial_number, string_to_bytes, Disk, Partition, PartitionTableType,
|
|
};
|
|
|
|
static DEFAULT_MAX_DISKS: usize = 8;
|
|
|
|
pub fn get_disk_details(disk_id: usize, disk_size: u64, disk_details: Option<&str>) -> Disk {
|
|
static RE_DETAILS: Lazy<Regex> = Lazy::new(|| {
|
|
Regex::new(r"(.*?)\r?\nDisk ID\s*:\s+(.*?)\r?\nType\s*:\s+(.*?)\r?\n").unwrap()
|
|
});
|
|
static RE_UUID: Lazy<Regex> = Lazy::new(|| {
|
|
Regex::new(r"^\{[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}\}$").unwrap()
|
|
});
|
|
let mut disk = Disk {
|
|
id: disk_id,
|
|
size: disk_size,
|
|
..Default::default()
|
|
};
|
|
|
|
// Get details
|
|
let details: String;
|
|
if let Some(details_str) = disk_details {
|
|
details = String::from(details_str);
|
|
} else {
|
|
let script = format!("select disk {disk_id}\r\ndetail disk");
|
|
details = run_script(&script);
|
|
};
|
|
|
|
// Parse details
|
|
for (_, [model, part_type, conn_type]) in
|
|
RE_DETAILS.captures_iter(&details).map(|c| c.extract())
|
|
{
|
|
disk.model = String::from(model);
|
|
disk.conn_type = String::from(conn_type);
|
|
if RE_UUID.is_match(part_type) {
|
|
disk.part_type = PartitionTableType::Guid;
|
|
} else {
|
|
disk.part_type = PartitionTableType::Legacy;
|
|
}
|
|
disk.serial = get_disk_serial_number(disk_id);
|
|
}
|
|
|
|
// Done
|
|
disk
|
|
}
|
|
|
|
pub fn get_partition_details(
|
|
disk_id: usize,
|
|
disk_details: Option<&str>,
|
|
part_details: Option<&str>,
|
|
) -> Vec<Partition> {
|
|
static RE_LIS: Lazy<Regex> =
|
|
Lazy::new(|| Regex::new(r"Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+B)").unwrap());
|
|
let mut parts = Vec::new();
|
|
|
|
// List partition
|
|
let contents: String;
|
|
if let Some(details) = disk_details {
|
|
contents = String::from(details);
|
|
} else {
|
|
let script = format!("select disk {disk_id}\r\nlist partition");
|
|
contents = run_script(&script);
|
|
};
|
|
for (_, [number, size]) in RE_LIS.captures_iter(&contents).map(|c| c.extract()) {
|
|
let part_num = number.parse().unwrap();
|
|
if part_num != 0 {
|
|
// part_num == 0 is reserved for extended partition "containers" so we can exclude them
|
|
let part = Partition {
|
|
id: part_num,
|
|
size: string_to_bytes(size),
|
|
..Default::default()
|
|
};
|
|
parts.push(part);
|
|
}
|
|
}
|
|
|
|
// Detail parititon
|
|
let mut script = vec![format!("select disk {}", disk_id)];
|
|
for part in &parts {
|
|
if part_details.is_some() {
|
|
// Currently only used by tests
|
|
break;
|
|
}
|
|
script.push(format!(
|
|
// Remove/Assign included to ensure all (accessible) volumes have letters
|
|
"\r\nselect partition {}\r\nremove noerr\r\nassign noerr\r\ndetail partition",
|
|
part.id
|
|
));
|
|
}
|
|
let part_contents: String;
|
|
if let Some(details) = part_details {
|
|
part_contents = String::from(details);
|
|
} else {
|
|
part_contents = run_script(script.join("\r\n").as_str());
|
|
};
|
|
parse_partition_details(&mut parts, &part_contents);
|
|
|
|
// Done
|
|
parts
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn build_dest_format_script(disk_id: usize, part_type: &PartitionTableType) -> String {
|
|
let disk_id = format!("{disk_id}");
|
|
let mut script = vec!["select disk {disk_id}", "clean"];
|
|
match part_type {
|
|
PartitionTableType::Guid => {
|
|
script.push("convert gpt");
|
|
script.push("create partition efi size=260");
|
|
script.push("format fs=fat32 quick label=ESP");
|
|
script.push("create partition msr size=16");
|
|
}
|
|
PartitionTableType::Legacy => {
|
|
script.push("create partition primary size=100");
|
|
script.push("active");
|
|
script.push("format fs=ntfs quick label=System");
|
|
}
|
|
}
|
|
script.join("\r\n").replace("{disk_id}", &disk_id)
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn build_get_disk_script(disk_nums: Option<Vec<&str>>) -> String {
|
|
let capacity = DEFAULT_MAX_DISKS * 3 + 1;
|
|
let script: String;
|
|
|
|
// Get disk and partition details
|
|
if let Some(disks) = disk_nums {
|
|
// (Slower case)
|
|
let mut script_parts = Vec::with_capacity(capacity);
|
|
|
|
// Get list of disks
|
|
script_parts.push(String::from("list disk"));
|
|
|
|
// Add disks from provided list
|
|
for num in disks {
|
|
script_parts.push(format!("select disk {num}"));
|
|
script_parts.push(String::from("detail disk"));
|
|
script_parts.push(String::from("list partition"));
|
|
}
|
|
|
|
// Done
|
|
script = script_parts.join("\n");
|
|
} else {
|
|
// (Best case)
|
|
let mut script_parts = Vec::with_capacity(capacity);
|
|
|
|
// Get list of disks
|
|
script_parts.push("list disk");
|
|
|
|
// Assuming first disk number is zero
|
|
script_parts.push("select disk 0");
|
|
script_parts.push("detail disk");
|
|
script_parts.push("list partition");
|
|
|
|
// Limit to 8 disks (if there's more the manual "worst" case will be used)
|
|
let mut i = 0;
|
|
while i < 8 {
|
|
script_parts.push("select disk next");
|
|
script_parts.push("detail disk");
|
|
script_parts.push("list partition");
|
|
i += 1;
|
|
}
|
|
|
|
// Done
|
|
script = script_parts.join("\n");
|
|
}
|
|
|
|
script
|
|
}
|
|
|
|
pub fn get_disks() -> Vec<Disk> {
|
|
static RE_DIS_DET: Lazy<Regex> = Lazy::new(|| {
|
|
Regex::new(r"(?s)Disk (\d+) is now the selected disk.*?\r?\n\s*\r?\n(.*)").unwrap()
|
|
});
|
|
static RE_DIS_LIS: Lazy<Regex> =
|
|
Lazy::new(|| Regex::new(r"Disk\s+(\d+)\s+(\w+)\s+(\d+\s+\w+B)").unwrap());
|
|
let mut contents: String;
|
|
let mut output;
|
|
let mut script: String;
|
|
|
|
// Run diskpart and check result
|
|
script = build_get_disk_script(None);
|
|
output = run_script_raw(&script);
|
|
contents = String::from_utf8_lossy(&output.stdout).to_string();
|
|
if let Some(return_code) = output.status.code() {
|
|
let disk_nums = parse_disk_numbers(&contents);
|
|
if return_code != 0 && !disk_nums.is_empty() {
|
|
// The base assumptions were correct! skipping fallback method
|
|
//
|
|
// Since the return_code was not zero, and at least one disk was detected, that
|
|
// means that the number of disks was less than DEFAULT_MAX_DISKS
|
|
//
|
|
// If the number of disks is exactly DEFAULT_MAX_DISKS then we may be omitting some
|
|
// disk(s) from the list so a manual check is needed
|
|
// This only adds one additional diskpart call, but avoiding it is preferred
|
|
//
|
|
// If zero disks were detected then either the first disk number wasn't zero or no
|
|
// disks were actually connected
|
|
// (this could be due to missing drivers but we're leaving that issue to the user)
|
|
info!("Disk assumptions correct!");
|
|
} else {
|
|
// The base assumptions were wrong so we need to check which disk numbers are present
|
|
warn!("Disk assumptions incorrect, falling back to manual method.");
|
|
script = build_get_disk_script(Some(disk_nums));
|
|
output = run_script_raw(&script);
|
|
contents = String::from_utf8_lossy(&output.stdout).to_string();
|
|
}
|
|
}
|
|
|
|
// Split Diskpart output contents into sections
|
|
let mut dp_sections = split_diskpart_disk_output(&contents);
|
|
|
|
// Build Disk structs
|
|
// NOTE: A HashMap is used because it's possible to have gaps in the list of disks
|
|
// i.e. 0, 1, 3, 4
|
|
// For instance, this can happen if a drive is disconnected after startup
|
|
let mut disks_map: HashMap<&str, Disk> = HashMap::with_capacity(DEFAULT_MAX_DISKS);
|
|
for (_, [number, _status, size]) in RE_DIS_LIS
|
|
.captures_iter(dp_sections.remove(0)) // This is the "list disk" section
|
|
.map(|c| c.extract())
|
|
{
|
|
disks_map.insert(
|
|
number,
|
|
Disk {
|
|
id: number.parse().unwrap(),
|
|
size: string_to_bytes(size),
|
|
..Default::default()
|
|
},
|
|
);
|
|
}
|
|
|
|
// Add Disk details
|
|
let mut disks_raw: Vec<Disk> = Vec::with_capacity(DEFAULT_MAX_DISKS);
|
|
for section in dp_sections {
|
|
for (_, [id, details]) in RE_DIS_DET.captures_iter(section).map(|c| c.extract()) {
|
|
if let Some(disk) = disks_map.remove(id) {
|
|
// We remove the disk from the HashMap because we're moving it to the Vec
|
|
let mut disk = get_disk_details(disk.id, disk.size, Some(details));
|
|
disk.parts = get_partition_details(disk.id, Some(details), None);
|
|
disk.generate_descriptions();
|
|
disks_raw.push(disk);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Return list of Disks
|
|
disks_raw
|
|
}
|
|
|
|
pub fn parse_disk_numbers(contents: &str) -> Vec<&str> {
|
|
//Disk 0 is now the selected disk.
|
|
//
|
|
//Red Hat VirtIO SCSI Disk Device
|
|
//Disk ID: {E9CE8DFA-46B2-43C1-99BB-850C661CEE6B}
|
|
static RE: Lazy<Regex> =
|
|
Lazy::new(|| Regex::new(r"\s+Disk\s+(\d+).*\n.*\n.*\nDisk ID:").unwrap());
|
|
|
|
let mut disk_nums = Vec::new();
|
|
for (_, [number]) in RE.captures_iter(contents).map(|c| c.extract()) {
|
|
disk_nums.push(number);
|
|
}
|
|
disk_nums
|
|
}
|
|
|
|
pub fn parse_partition_details(parts: &mut [Partition], contents: &str) {
|
|
static RE_PAR: Lazy<Regex> = Lazy::new(|| {
|
|
Regex::new(
|
|
r"Partition (\d+)\r?\nType\s*: (\S+)(\r?\n.*){5}\s*(Volume.*\r?\n.*\r?\n|There is no volume)(.*)",
|
|
)
|
|
.unwrap()
|
|
});
|
|
static RE_VOL: Lazy<Regex> = Lazy::new(|| {
|
|
// Volume ### Ltr Label Fs Type Size Status Info
|
|
// ---------- --- ----------- ----- ---------- ------- --------- --------
|
|
// * Volume 1 S ESP FAT32 Partition 100 MB Healthy Hidden
|
|
Regex::new(r"..Volume (\d.{2}) (.{3}) (.{11}) (.{5})").unwrap()
|
|
});
|
|
|
|
for (part_index, (_, [_part_id, part_type, _, _vol_header, vol_line])) in RE_PAR
|
|
.captures_iter(contents)
|
|
.map(|c| c.extract())
|
|
.enumerate()
|
|
{
|
|
if let Some(part) = parts.get_mut(part_index) {
|
|
// Partition info
|
|
part.part_type = String::from(part_type.trim());
|
|
|
|
// Volume info
|
|
for (_, [_id, letter, label, fs_type]) in
|
|
RE_VOL.captures_iter(vol_line).map(|c| c.extract())
|
|
{
|
|
part.label = String::from(label.trim());
|
|
part.letter = String::from(letter.trim());
|
|
part.fs_type = String::from(fs_type.trim());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn refresh_disk_info(disk: &Disk) -> Disk {
|
|
info!("Refresh disk info");
|
|
let mut disk = get_disk_details(disk.id, disk.size, None);
|
|
disk.parts = get_partition_details(disk.id, None, None);
|
|
disk.generate_descriptions();
|
|
disk
|
|
}
|
|
|
|
/// # Panics
|
|
///
|
|
/// Will panic if diskpart fails to run the script or if the script if malformed
|
|
pub fn run_script_raw(script: &str) -> Output {
|
|
info!("Running Diskpart: {:?}", &script);
|
|
let temp_dir = tempdir().expect("Failed to create temp dir");
|
|
let script_path = temp_dir.path().join("diskpart.script");
|
|
let mut script_file = File::create(&script_path).expect("Failed to create temp file");
|
|
script_file
|
|
.write_all(script.as_bytes())
|
|
.expect("Failed to write script to disk");
|
|
let output = Command::new("diskpart")
|
|
.args(["/s", format!("{}", script_path.display()).as_str()])
|
|
.stdout(Stdio::piped())
|
|
.output()
|
|
.expect("Failed to execute Diskpart script");
|
|
output
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn run_script(script: &str) -> String {
|
|
let output = run_script_raw(script);
|
|
String::from_utf8_lossy(&output.stdout).to_string()
|
|
}
|
|
|
|
pub fn split_diskpart_disk_output(contents: &str) -> Vec<&str> {
|
|
// NOTE: A simple split isn't helpful since we want to include the matching lines
|
|
static RE: Lazy<Regex> =
|
|
Lazy::new(|| Regex::new(r"Disk \d+ is now the selected disk").unwrap());
|
|
let mut sections = Vec::new();
|
|
let mut starts: Vec<usize> = vec![0];
|
|
let mut ends: Vec<usize> = Vec::new();
|
|
let _: Vec<_> = RE
|
|
.find_iter(contents)
|
|
.map(|m| {
|
|
ends.push(m.start() - 1);
|
|
starts.push(m.start());
|
|
})
|
|
.collect();
|
|
ends.push(contents.len());
|
|
let ranges: Vec<(&usize, &usize)> = starts.iter().zip(ends.iter()).collect();
|
|
for range in ranges {
|
|
let start = *range.0;
|
|
let end = *range.1;
|
|
sections.push(&contents[start..end]);
|
|
}
|
|
sections
|
|
}
|