1use regex::Regex;
2use serde_json::Value;
3use std::collections::HashMap;
4use std::str;
5
6use crate::squire;
7
8fn parse_size(size_str: &str) -> String {
18 let re = Regex::new(r"([\d.]+)([KMGTP]?)").unwrap();
19 if let Some(caps) = re.captures(size_str.trim()) {
20 let value: f64 = caps[1].parse().unwrap();
21 let unit = &caps[2];
22 let unit_multipliers = HashMap::from([
23 ("K", 2_f64.powi(10)),
24 ("M", 2_f64.powi(20)),
25 ("G", 2_f64.powi(30)),
26 ("T", 2_f64.powi(40)),
27 ("P", 2_f64.powi(50)),
28 ]);
29 let multiplier = unit_multipliers.get(unit).unwrap_or(&1.0);
30 return squire::util::size_converter((value * multiplier) as u64);
31 }
32 size_str.replace("K", " KB")
33 .replace("M", " MB")
34 .replace("G", " GB")
35 .replace("T", " TB")
36 .replace("P", " PB")
37}
38
39fn is_physical_disk(lib_path: &str, device_id: &str) -> bool {
50 let result = squire::util::run_command(lib_path, &["info", device_id], true);
51 let output = match result {
52 Ok(output) => output,
53 Err(_) => {
54 log::error!("Failed to get disk info");
55 return false;
56 }
57 };
58 for line in output.split("\n") {
59 if line.contains("Virtual:") && line.contains("Yes") {
60 return false;
61 }
62 }
63 true
64}
65
66fn linux_disks(lib_path: &str) -> Vec<HashMap<String, String>> {
76 let result = squire::util::run_command(lib_path, &["-o", "NAME,SIZE,TYPE,MODEL", "-d"], true);
77 let output = match result {
78 Ok(output) => output,
79 Err(_) => {
80 log::error!("Failed to get disk info");
81 return Vec::new();
82 }
83 };
84 let disks_skipped: Vec<&str> = output.lines().skip(1).collect();
86 let filtered_disks: Vec<&str> = disks_skipped.into_iter().filter(|&disk| !disk.contains("loop")).collect();
87 let mut disk_list = Vec::new();
88 for disk in filtered_disks {
89 let parts: Vec<&str> = disk.split_whitespace().collect();
91 if parts.len() >= 4 {
93 let disk_info = HashMap::from([
94 ("Name".to_string(), parts[0].to_string()),
95 ("Size".to_string(), parse_size(parts[1])),
96 ("Type".to_string(), parts[2].to_string()),
97 ("Model".to_string(), parts[3..].join(" ")),
98 ]);
99 disk_list.push(disk_info);
100 }
101 }
102 disk_list
103}
104
105fn darwin_disks(lib_path: &str) -> Vec<HashMap<String, String>> {
115 let result = squire::util::run_command(lib_path, &["list"], true);
116 let output = match result {
117 Ok(output) => output,
118 Err(_) => {
119 log::error!("Failed to get disk info");
120 return Vec::new();
121 }
122 };
123 let disks: Vec<&str> = output.lines().collect();
124 let disk_lines: Vec<&str> = disks
125 .into_iter()
126 .filter(|&line| line.starts_with("/dev/disk"))
127 .collect();
128 let mut disk_info = Vec::new();
129 for line in disk_lines {
130 let device_id = line
131 .split_whitespace()
132 .next()
133 .unwrap_or_default();
134 if !is_physical_disk(lib_path, device_id) {
135 continue;
136 }
137 let result = squire::util::run_command(lib_path, &["info", device_id], true);
138 let disk_info_output = match result {
139 Ok(output) => output,
140 Err(_) => {
141 log::error!("Failed to get disk info");
142 return Vec::new();
143 }
144 };
145 let info_lines: Vec<&str> = disk_info_output
146 .lines()
147 .collect();
148 let mut disk_data = HashMap::new();
149 for info_line in info_lines {
150 if info_line.contains("Device / Media Name:") {
151 disk_data.insert(
152 "Name".to_string(),
153 info_line
154 .split(":")
155 .nth(1)
156 .unwrap_or_default()
157 .trim()
158 .to_string()
159 );
160 }
161 if info_line.contains("Disk Size:") {
162 let size_info = info_line
163 .split(":")
164 .nth(1)
165 .unwrap_or_default()
166 .split("(")
167 .next()
168 .unwrap_or_default()
169 .trim()
170 .to_string();
171 disk_data.insert(
172 "Size".to_string(),
173 size_info
174 );
175 }
176 }
177 disk_data.insert(
178 "DeviceID".to_string(),
179 device_id.to_string()
180 );
181 disk_info.push(disk_data);
182 }
183 disk_info
184}
185
186fn reformat_windows(data: &mut HashMap<String, Value>) -> HashMap<String, String> {
196 let size = data.get("Size").unwrap().as_f64().unwrap();
197 let model = data.get("Model").unwrap().as_str().unwrap().to_string();
198 let mut reformatted_data = HashMap::new();
199 reformatted_data.insert("Size".to_string(), squire::util::size_converter(size as u64));
200 reformatted_data.insert("Name".to_string(), model);
201 reformatted_data.insert(
202 "DeviceID".to_string(),
203 data.get("DeviceID")
204 .unwrap_or(&Value::String("".to_string()))
205 .as_str()
206 .unwrap_or_default()
207 .to_string()
208 );
209 reformatted_data
210}
211
212fn windows_disks(lib_path: &str) -> Vec<HashMap<String, String>> {
222 let ps_command = "Get-CimInstance Win32_DiskDrive | Select-Object Caption, DeviceID, Model, Partitions, Size | ConvertTo-Json";
223 let result = squire::util::run_command(lib_path, &["-Command", ps_command], true);
224 let output = match result {
225 Ok(output) => output,
226 Err(_) => {
227 log::error!("Failed to get disk info");
228 return Vec::new();
229 }
230 };
231 let disks_info: Value = serde_json::from_str(&output).unwrap();
232 let mut disk_info = Vec::new();
233 if let Some(disks) = disks_info.as_array() {
234 for disk in disks {
235 let mut disk_map: HashMap<String, Value> = serde_json::from_value(disk.clone()).unwrap();
236 disk_info.push(reformat_windows(&mut disk_map));
237 }
238 } else {
239 let mut disk_map: HashMap<String, Value> = serde_json::from_value(disks_info).unwrap();
240 disk_info.push(reformat_windows(&mut disk_map));
241 }
242 disk_info
243}
244
245pub fn get_all_disks() -> Vec<HashMap<String, String>> {
251 let operating_system = std::env::consts::OS;
252 match operating_system {
253 "windows" => windows_disks("C:\\Program Files\\PowerShell\\7\\pwsh.exe"),
254 "macos" => darwin_disks("/usr/sbin/diskutil"),
255 "linux" => linux_disks("/usr/bin/lsblk"),
256 _ => {
257 log::error!("Unsupported operating system");
258 Vec::new()
259 }
260 }
261}