Move set/refresh diskpart logic to Disk

This commit is contained in:
2Shirt 2025-04-22 16:25:30 -07:00
parent 88f1234dd0
commit dc49006af8
Signed by: 2Shirt
GPG key ID: 152FAC923B0E132C
4 changed files with 70 additions and 49 deletions

View file

@ -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<Partition>,
pub parts_description: Vec<String>,
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<Partition> {
vec![
Partition {

View file

@ -32,7 +32,7 @@ use crate::system::disk::{
static DEFAULT_MAX_DISKS: usize = 8;
struct RegexList {
pub struct RegexList {
detail_all_disks: OnceLock<Regex>,
detail_disk: OnceLock<Regex>,
detail_partition: OnceLock<Regex>,
@ -41,7 +41,7 @@ struct RegexList {
list_partition: OnceLock<Regex>,
list_volumes: OnceLock<Regex>,
split_all_disks: OnceLock<Regex>,
re_uuid: OnceLock<Regex>,
uuid: OnceLock<Regex>,
}
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<Disk> {
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

View file

@ -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 => {

View file

@ -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),