diff --git a/core/src/system/disk.rs b/core/src/system/disk.rs index bcf01c8..f456de6 100644 --- a/core/src/system/disk.rs +++ b/core/src/system/disk.rs @@ -25,7 +25,7 @@ use tracing::info; use regex::Regex; -use crate::system::diskpart; +use crate::system::diskpart::{self, REGEXES}; #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct Disk { @@ -36,7 +36,6 @@ pub struct Disk { pub part_type: PartitionTableType, pub parts: Vec, pub parts_description: Vec, - pub sector_size: usize, pub serial: String, pub size: u64, // In bytes } @@ -86,6 +85,44 @@ impl Disk { pub fn num_parts(&self) -> usize { self.parts.len() } + + pub fn refresh_disk_info(&mut self, details_str: Option<&str>) { + let re_details = REGEXES.detail_disk(); + let re_uuid = REGEXES.uuid(); + if cfg!(windows) { + info!("Refresh disk info via Diskpart"); + // Get details + let details: String; + if let Some(s) = details_str { + details = String::from(s); + } else { + let script = format!("select disk {}\r\ndetail disk", self.id); + details = diskpart::run_script(&script); + }; + + // Parse details + for (_, [model, part_type, conn_type]) in + re_details.captures_iter(&details).map(|c| c.extract()) + { + self.model = String::from(model); + self.conn_type = String::from(conn_type); + if re_uuid.is_match(part_type) { + self.part_type = PartitionTableType::Guid; + } else { + self.part_type = PartitionTableType::Legacy; + } + self.serial = get_disk_serial_number(self.id); + } + + // Partition details + self.parts = diskpart::get_partitions(self.id, None, None); + self.generate_descriptions(); + } else { + info!("Refresh fake disk info"); + self.parts = refresh_fake_disk_info(); + self.generate_descriptions(); + } + } } impl fmt::Display for Disk { @@ -106,6 +143,24 @@ impl fmt::Display for Disk { } } +impl Partition { + pub fn set_details(&mut self, part_type: &str, vol_line: &str) { + let re_list_volume = REGEXES.list_volumes(); + + // Partition info + self.part_type = String::from(part_type.trim()); + + // Volume info + for (_, [_id, letter, label, fs_type]) in + re_list_volume.captures_iter(vol_line).map(|c| c.extract()) + { + self.label = String::from(label.trim()); + self.letter = String::from(letter.trim()); + self.fs_type = String::from(fs_type.trim()); + } + } +} + impl fmt::Display for Partition { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { let mut s: String; @@ -315,20 +370,6 @@ pub fn get_disk_serial_number(id: usize) -> String { serial } -pub fn refresh_disk_info(disk: &mut Disk) -> Disk { - let mut new_disk: Disk; - if cfg!(windows) { - info!("Refresh disk via Diskpart"); - new_disk = diskpart::refresh_disk_info(disk); - } else { - info!("Refresh fake disk"); - new_disk = disk.clone(); - new_disk.parts = refresh_fake_disk_info(); - new_disk.generate_descriptions(); - } - new_disk -} - fn refresh_fake_disk_info() -> Vec { vec![ Partition { diff --git a/core/src/system/diskpart.rs b/core/src/system/diskpart.rs index 7a62bc5..7eb19ee 100644 --- a/core/src/system/diskpart.rs +++ b/core/src/system/diskpart.rs @@ -32,7 +32,7 @@ use crate::system::disk::{ static DEFAULT_MAX_DISKS: usize = 8; -struct RegexList { +pub struct RegexList { detail_all_disks: OnceLock, detail_disk: OnceLock, detail_partition: OnceLock, @@ -41,7 +41,7 @@ struct RegexList { list_partition: OnceLock, list_volumes: OnceLock, split_all_disks: OnceLock, - re_uuid: OnceLock, + uuid: OnceLock, } impl RegexList { @@ -92,14 +92,14 @@ impl RegexList { } pub fn uuid(&self) -> &Regex { - self.re_uuid.get_or_init(|| { + self.uuid.get_or_init(|| { Regex::new(r"^\{[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}\}$") .unwrap() }) } } -static REGEXES: RegexList = RegexList { +pub static REGEXES: RegexList = RegexList { detail_all_disks: OnceLock::new(), detail_disk: OnceLock::new(), detail_partition: OnceLock::new(), @@ -108,11 +108,11 @@ static REGEXES: RegexList = RegexList { list_partition: OnceLock::new(), list_volumes: OnceLock::new(), split_all_disks: OnceLock::new(), - re_uuid: OnceLock::new(), + uuid: OnceLock::new(), }; pub fn get_disk_details(disk_id: usize, disk_size: u64, disk_details: Option<&str>) -> Disk { - let re_details = REGEXES.detail_disk(); + let detail_disk = REGEXES.detail_disk(); let re_uuid = REGEXES.uuid(); let mut disk = Disk { id: disk_id, @@ -131,7 +131,7 @@ pub fn get_disk_details(disk_id: usize, disk_size: u64, disk_details: Option<&st // Parse details for (_, [model, part_type, conn_type]) in - re_details.captures_iter(&details).map(|c| c.extract()) + detail_disk.captures_iter(&details).map(|c| c.extract()) { disk.model = String::from(model); disk.conn_type = String::from(conn_type); @@ -147,7 +147,7 @@ pub fn get_disk_details(disk_id: usize, disk_size: u64, disk_details: Option<&st disk } -pub fn get_partition_details( +pub fn get_partitions( disk_id: usize, disk_details: Option<&str>, part_details: Option<&str>, @@ -338,7 +338,7 @@ pub fn get_disks() -> Vec { 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.parts = get_partitions(disk.id, Some(details), None); disk.generate_descriptions(); disks_raw.push(disk); } @@ -364,7 +364,6 @@ pub fn parse_disk_numbers(contents: &str) -> Vec<&str> { pub fn parse_partition_details(parts: &mut [Partition], contents: &str) { let detail_partition = REGEXES.detail_partition(); - let list_volume = REGEXES.list_volumes(); for (part_index, (_, [_part_id, part_type, _, _vol_header, vol_line])) in detail_partition .captures_iter(contents) @@ -372,29 +371,11 @@ pub fn parse_partition_details(parts: &mut [Partition], contents: &str) { .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 - list_volume.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()); - } + part.set_details(part_type, vol_line); } } } -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 diff --git a/core/src/tasks.rs b/core/src/tasks.rs index a548e34..a6c212f 100644 --- a/core/src/tasks.rs +++ b/core/src/tasks.rs @@ -236,8 +236,7 @@ impl Tasks { let disk_list_arc = self.disk_list.clone(); self.cur_handle = Some(thread::spawn(move || { let mut disks = disk_list_arc.lock().unwrap(); - let old_disk = &mut disks[index]; - disks[index] = disk::refresh_disk_info(old_disk); + disks[index].refresh_disk_info(None); })); } TaskType::UpdateDiskList => { diff --git a/core/src/tests/mod.rs b/core/src/tests/mod.rs index 6750d5d..c43759c 100644 --- a/core/src/tests/mod.rs +++ b/core/src/tests/mod.rs @@ -37,7 +37,7 @@ mod diskpart { } #[test] - fn get_partition_details() { + fn get_partitions() { // Left let partition_1 = system::disk::Partition { id: 1, @@ -77,7 +77,7 @@ mod diskpart { id: 4, ..Default::default() }); - disk.parts = system::diskpart::get_partition_details( + disk.parts = system::diskpart::get_partitions( 2, Some(sample_output::DETAIL_DISK_GPT), Some(sample_output::SELECT_PART_DETAIL_PARTS),