use serde::Deserialize;
use serde_json::Value;
use std::process::Command;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DriveInfo {
pub name: String,
pub path: String,
pub model: Option<String>,
pub serial: Option<String>,
pub size_bytes: u64,
pub removable: bool,
pub read_only: bool,
pub transport: Option<String>,
pub mountpoints: Vec<String>,
pub partitions: Vec<PartitionInfo>,
}
impl DriveInfo {
pub fn display_name(&self) -> String {
match self
.model
.as_deref()
.filter(|model| !model.trim().is_empty())
{
Some(model) => format!("{} ({})", model.trim(), self.path),
None => self.path.clone(),
}
}
pub fn human_size(&self) -> String {
human_size(self.size_bytes)
}
pub fn safety_badge(&self) -> &'static str {
if self.read_only {
"read-only"
} else if self.transport.as_deref() == Some("usb") {
"usb"
} else if self.removable {
"removable"
} else {
"hotplug"
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PartitionInfo {
pub name: String,
pub path: String,
pub size_bytes: u64,
pub mountpoints: Vec<String>,
}
impl PartitionInfo {
pub fn human_size(&self) -> String {
human_size(self.size_bytes)
}
}
pub trait DriveScanner {
fn scan(&self) -> Result<Vec<DriveInfo>, DriveScanError>;
}
#[derive(Debug, Default, Clone, Copy)]
pub struct LsblkScanner;
impl DriveScanner for LsblkScanner {
fn scan(&self) -> Result<Vec<DriveInfo>, DriveScanError> {
let output = Command::new("lsblk")
.args([
"--json",
"--bytes",
"--output",
"NAME,PATH,MODEL,SERIAL,SIZE,RM,RO,TYPE,MOUNTPOINTS,TRAN,HOTPLUG",
])
.output()?;
if !output.status.success() {
return Err(DriveScanError::LsblkFailed(
String::from_utf8_lossy(&output.stderr).trim().to_owned(),
));
}
parse_lsblk_json(&output.stdout)
}
}
#[derive(Debug, thiserror::Error)]
pub enum DriveScanError {
#[error("failed to run lsblk: {0}")]
Command(#[from] std::io::Error),
#[error("failed to parse lsblk JSON: {0}")]
Json(#[from] serde_json::Error),
#[error("lsblk failed: {0}")]
LsblkFailed(String),
}
pub fn parse_lsblk_json(json: impl AsRef<[u8]>) -> Result<Vec<DriveInfo>, DriveScanError> {
let lsblk: LsblkOutput = serde_json::from_slice(json.as_ref())?;
let mut drives = lsblk
.blockdevices
.into_iter()
.filter_map(device_to_drive)
.collect::<Vec<_>>();
drives.sort_by(|left, right| left.path.cmp(&right.path));
Ok(drives)
}
fn device_to_drive(device: LsblkDevice) -> Option<DriveInfo> {
if device.kind.as_deref() != Some("disk") {
return None;
}
let transport = clean(device.tran);
let is_usb = transport.as_deref() == Some("usb");
let removable = device.rm || device.hotplug || is_usb;
if !removable {
return None;
}
let path = clean(device.path)?;
let name = clean(device.name).unwrap_or_else(|| path.trim_start_matches("/dev/").to_owned());
let mut mountpoints = compact_mountpoints(device.mountpoints);
let partitions = device
.children
.unwrap_or_default()
.into_iter()
.filter_map(|child| {
let child_path = clean(child.path)?;
let child_name = clean(child.name).unwrap_or_else(|| child_path.clone());
let child_mountpoints = compact_mountpoints(child.mountpoints);
mountpoints.extend(child_mountpoints.iter().cloned());
Some(PartitionInfo {
name: child_name,
path: child_path,
size_bytes: child.size.unwrap_or(0),
mountpoints: child_mountpoints,
})
})
.collect::<Vec<_>>();
mountpoints.sort();
mountpoints.dedup();
Some(DriveInfo {
name,
path,
model: clean(device.model),
serial: clean(device.serial),
size_bytes: device.size.unwrap_or(0),
removable,
read_only: device.ro,
transport,
mountpoints,
partitions,
})
}
fn compact_mountpoints(mountpoints: Vec<Option<String>>) -> Vec<String> {
mountpoints
.into_iter()
.flatten()
.filter_map(|mountpoint| clean(Some(mountpoint)))
.collect()
}
fn clean(value: Option<String>) -> Option<String> {
let value = value?.trim().to_owned();
(!value.is_empty()).then_some(value)
}
pub fn human_size(bytes: u64) -> String {
const UNITS: [&str; 5] = ["B", "KiB", "MiB", "GiB", "TiB"];
let mut value = bytes as f64;
let mut unit = 0;
while value >= 1024.0 && unit < UNITS.len() - 1 {
value /= 1024.0;
unit += 1;
}
if unit == 0 {
format!("{} {}", bytes, UNITS[unit])
} else {
format!("{value:.1} {}", UNITS[unit])
}
}
#[derive(Debug, Deserialize)]
struct LsblkOutput {
blockdevices: Vec<LsblkDevice>,
}
#[derive(Debug, Deserialize)]
struct LsblkDevice {
name: Option<String>,
path: Option<String>,
model: Option<String>,
serial: Option<String>,
#[serde(default, deserialize_with = "deserialize_optional_u64")]
size: Option<u64>,
#[serde(default, deserialize_with = "deserialize_boolish")]
rm: bool,
#[serde(default, deserialize_with = "deserialize_boolish")]
ro: bool,
#[serde(rename = "type")]
kind: Option<String>,
#[serde(default, deserialize_with = "deserialize_mountpoints")]
mountpoints: Vec<Option<String>>,
tran: Option<String>,
#[serde(default, deserialize_with = "deserialize_boolish")]
hotplug: bool,
#[serde(default)]
children: Option<Vec<LsblkDevice>>,
}
fn deserialize_optional_u64<'de, D>(deserializer: D) -> Result<Option<u64>, D::Error>
where
D: serde::Deserializer<'de>,
{
let value = Option::<Value>::deserialize(deserializer)?;
Ok(match value {
Some(Value::Number(number)) => number.as_u64(),
Some(Value::String(text)) => text.trim().parse().ok(),
_ => None,
})
}
fn deserialize_boolish<'de, D>(deserializer: D) -> Result<bool, D::Error>
where
D: serde::Deserializer<'de>,
{
let value = Option::<Value>::deserialize(deserializer)?;
Ok(match value {
Some(Value::Bool(value)) => value,
Some(Value::Number(number)) => number.as_u64().unwrap_or(0) != 0,
Some(Value::String(text)) => matches!(
text.trim().to_ascii_lowercase().as_str(),
"1" | "true" | "yes" | "y"
),
_ => false,
})
}
fn deserialize_mountpoints<'de, D>(deserializer: D) -> Result<Vec<Option<String>>, D::Error>
where
D: serde::Deserializer<'de>,
{
let value = Option::<Value>::deserialize(deserializer)?;
Ok(match value {
Some(Value::Array(values)) => values
.into_iter()
.map(|value| match value {
Value::String(text) => Some(text),
_ => None,
})
.collect(),
Some(Value::String(text)) => vec![Some(text)],
_ => Vec::new(),
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_only_removable_drives() {
let json = br#"
{
"blockdevices": [
{"name":"nvme0n1","path":"/dev/nvme0n1","model":"Internal","size":512110190592,"rm":false,"ro":false,"type":"disk","mountpoints":[null],"tran":"nvme","hotplug":false},
{"name":"sdb","path":"/dev/sdb","model":"USB Stick","serial":"ABC","size":"16008609792","rm":true,"ro":false,"type":"disk","mountpoints":[null],"tran":"usb","hotplug":true,
"children":[{"name":"sdb1","path":"/dev/sdb1","size":16000000000,"rm":true,"ro":false,"type":"part","mountpoints":["/run/media/usb"]}]}
]
}
"#;
let drives = parse_lsblk_json(json).expect("parse lsblk");
assert_eq!(drives.len(), 1);
assert_eq!(drives[0].path, "/dev/sdb");
assert_eq!(drives[0].model.as_deref(), Some("USB Stick"));
assert_eq!(drives[0].mountpoints, vec!["/run/media/usb"]);
assert_eq!(drives[0].partitions[0].path, "/dev/sdb1");
}
#[test]
fn accepts_numeric_bool_fields() {
let json = br#"
{
"blockdevices": [
{"name":"sdc","path":"/dev/sdc","model":"Reader","size":"1024","rm":0,"ro":0,"type":"disk","mountpoints":null,"tran":"usb","hotplug":1}
]
}
"#;
let drives = parse_lsblk_json(json).expect("parse lsblk");
assert_eq!(drives.len(), 1);
assert_eq!(drives[0].size_bytes, 1024);
assert_eq!(drives[0].safety_badge(), "usb");
}
#[test]
fn formats_human_sizes() {
assert_eq!(human_size(0), "0 B");
assert_eq!(human_size(1024), "1.0 KiB");
assert_eq!(human_size(4 * 1024 * 1024 * 1024), "4.0 GiB");
}
}