sysmonk/legacy/
disks.rs

1use regex::Regex;
2use serde_json::Value;
3use std::collections::HashMap;
4use std::str;
5
6use crate::squire;
7
8/// Function to parse size string for Linux.
9///
10/// # Arguments
11///
12/// * `size_str` - The size string to parse
13///
14/// # Returns
15///
16/// A `String` containing the parsed size string.
17fn 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
39/// Function to check if a disk is physical/virtual for macOS.
40///
41/// # Arguments
42///
43/// * `lib_path` - The path to the library
44/// * `device_id` - The device ID
45///
46/// # Returns
47///
48/// A `bool` indicating if the disk is physical.
49fn 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
66/// Function to get disk information on Linux.
67///
68/// # Arguments
69///
70/// * `lib_path` - The path to the library used to get disks' information.
71///
72/// # Returns
73///
74/// A `Vec` of `HashMap` containing the disk information.
75fn 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    // Skip the header line
85    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        // Split the disk info by whitespace and collect each part
90        let parts: Vec<&str> = disk.split_whitespace().collect();
91        // Ensure the disk info has at least 4 parts (NAME, SIZE, TYPE, MODEL)
92        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
105/// Function to get disk information on macOS.
106///
107/// # Arguments
108///
109/// * `lib_path` - The path to the library used to get disks' information.
110///
111/// # Returns
112///
113/// A `Vec` of `HashMap` containing the disk information.
114fn 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
186/// Function to reformat disk information on Windows.
187///
188/// # Arguments
189///
190/// * `data` - A mutable reference to the disk information.
191///
192/// # Returns
193///
194/// A `HashMap` containing the reformatted disk information.
195fn 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
212/// Function to get disk information on Windows.
213///
214/// # Arguments
215///
216/// * `lib_path` - The path to the library used to get disks' information.
217///
218/// # Returns
219///
220/// A `Vec` of `HashMap` containing the disk information.
221fn 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
245/// Function to get all disks' information.
246///
247/// # Returns
248///
249/// A `Vec` of `HashMap` containing the disk information.
250pub 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}