diff --git a/src/system/diskpart.rs b/src/system/diskpart.rs index b2f7d9d..a4bee11 100644 --- a/src/system/diskpart.rs +++ b/src/system/diskpart.rs @@ -31,20 +31,25 @@ use crate::system::disk::{ static DEFAULT_MAX_DISKS: usize = 8; -pub fn add_disk_details(disk: &mut Disk, disk_details: Option<&str>) { +pub fn get_disk_details(disk_id: &str, disk_size: u64, disk_details: Option<&str>) -> Disk { static RE_DETAILS: Lazy = Lazy::new(|| { Regex::new(r"(.*?)\r?\nDisk ID\s*:\s+(.*?)\r?\nType\s*:\s+(.*?)\r?\n").unwrap() }); static RE_UUID: Lazy = 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: String::from(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 {}\r\ndetail disk", disk.get_id()); + let script = format!("select disk {}\r\ndetail disk", disk_id); details = run_script(&script); }; @@ -59,8 +64,11 @@ pub fn add_disk_details(disk: &mut Disk, disk_details: Option<&str>) { } else { disk.part_type = PartitionTableType::Legacy; } - disk.serial = get_disk_serial_number(&disk.get_id()); + disk.serial = get_disk_serial_number(disk_id); } + + // Done + disk } pub fn add_partition_details( @@ -249,9 +257,9 @@ pub fn get_disks() -> Vec { let mut disks_raw: Vec = 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(mut disk) = disks_map.remove(id) { + if let Some(disk) = disks_map.remove(id) { // We remove the disk from the HashMap because we're moving it to the Vec - add_disk_details(&mut disk, Some(details)); + let mut disk = get_disk_details(disk.id.as_str(), disk.size, Some(details)); add_partition_details(&mut disk, Some(details), None); disk.generate_descriptions(); disks_raw.push(disk); @@ -323,13 +331,11 @@ pub fn parse_partition_details(disk: &mut Disk, contents: &str) { } } -pub fn refresh_disk_info(disk: &mut Disk) { +pub fn refresh_disk_info(disk: &Disk) { // TODO: Needs refactor - assuming add_ functions are replaced with get_ variants info!("Refresh disk info"); - disk.parts.clear(); - disk.parts_description.clear(); - add_disk_details(disk, None); - add_partition_details(disk, None, None); + let mut disk = get_disk_details(disk.get_id(), disk.size, None); + add_partition_details(&mut disk, None, None); disk.generate_descriptions(); } diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 786bf81..5b03d18 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -22,9 +22,7 @@ mod diskpart { #[test] fn add_disk_details_gpt() { - let mut disk = system::disk::Disk::default(); - - system::diskpart::add_disk_details(&mut disk, Some(sample_output::DETAIL_DISK_GPT)); + let disk = system::diskpart::get_disk_details("", 0, Some(sample_output::DETAIL_DISK_GPT)); assert_eq!(disk.model, "Red Hat VirtIO SCSI Disk Device"); assert_eq!(disk.part_type, system::disk::PartitionTableType::Guid); assert_eq!(disk.conn_type, "SCSI"); @@ -32,9 +30,7 @@ mod diskpart { #[test] fn add_disk_details_mbr() { - let mut disk = system::disk::Disk::default(); - - system::diskpart::add_disk_details(&mut disk, Some(sample_output::DETAIL_DISK_MBR)); + let disk = system::diskpart::get_disk_details("", 0, Some(sample_output::DETAIL_DISK_MBR)); assert_eq!(disk.model, "Red Hat VirtIO SCSI Disk Device"); assert_eq!(disk.part_type, system::disk::PartitionTableType::Legacy); assert_eq!(disk.conn_type, "SCSI");