diskr/analyzer/

disk_health.rs

use serde::Serialize;
use std::path::Path;
use std::process::Command;
use sysinfo::Disks;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub enum HealthStatus {
    Good,
    Warning,
    Critical,
    Unknown,
}

#[derive(Debug, Clone, Serialize)]
pub struct DiskHealth {
    pub device: String,
    pub model: String,
    pub kind: String,
    pub total_bytes: u64,
    pub available_bytes: u64,
    pub temperature_c: Option<i32>,
    pub power_on_hours: Option<u64>,
    pub reallocated_sectors: Option<u64>,
    pub wear_level_percent: Option<u8>,
    pub status: HealthStatus,
    pub smart_available: bool,
}

pub fn list_disks() -> Vec<DiskHealth> {
    let disks = Disks::new_with_refreshed_list();
    disks
        .iter()
        .map(|d| {
            let device = d.name().to_string_lossy().to_string();
            let phys = physical_device(&device);
            let smart = read_smart(&phys);
            let smart_available = smart.is_some();

            let model = smart
                .as_ref()
                .and_then(|s| s.model.clone())
                .or_else(|| read_model_sysfs(&phys))
                .unwrap_or_else(|| "Unknown model".to_string());

            let mut health = DiskHealth {
                device,
                model,
                kind: format!("{:?}", d.kind()),
                total_bytes: d.total_space(),
                available_bytes: d.available_space(),
                temperature_c: smart.as_ref().and_then(|s| s.temperature_c),
                power_on_hours: smart.as_ref().and_then(|s| s.power_on_hours),
                reallocated_sectors: smart.as_ref().and_then(|s| s.reallocated_sectors),
                wear_level_percent: smart.as_ref().and_then(|s| s.wear_level_percent),
                status: HealthStatus::Unknown,
                smart_available,
            };
            health.status = evaluate(&health);
            health
        })
        .collect()
}

fn physical_device(device: &str) -> String {
    let path = Path::new(device);
    let dir = path.parent().and_then(|p| p.to_str()).unwrap_or("/dev");
    let name = match path.file_name().and_then(|n| n.to_str()) {
        Some(n) => n,
        None => return device.to_string(),
    };

    if name.starts_with("nvme") || name.starts_with("mmcblk") {
        if let Some(pidx) = name.rfind('p') {
            let tail = &name[pidx + 1..];
            if !tail.is_empty() && tail.chars().all(|c| c.is_ascii_digit()) {
                return format!("{dir}/{}", &name[..pidx]);
            }
        }
        return format!("{dir}/{name}");
    }

    // sda1 -> sda, vdb2 -> vdb, hdc3 -> hdc, xvda1 -> xvda
    if name.starts_with("sd") || name.starts_with("hd") || name.starts_with("vd") || name.starts_with("xvd") {
        let trimmed = name.trim_end_matches(|c: char| c.is_ascii_digit());
        if !trimmed.is_empty() {
            return format!("{dir}/{trimmed}");
        }
    }

    device.to_string()
}

fn read_model_sysfs(phys_device: &str) -> Option<String> {
    let base_name = Path::new(phys_device).file_name()?.to_str()?;
    let candidates = [
        format!("/sys/block/{base_name}/device/model"),
        format!("/sys/block/{base_name}/device/name"),
        format!("/sys/class/nvme/{base_name}/model"),
    ];
    for c in candidates {
        if let Ok(text) = std::fs::read_to_string(&c) {
            let trimmed = text.trim();
            if !trimmed.is_empty() {
                return Some(trimmed.to_string());
            }
        }
    }
    None
}

struct SmartInfo {
    model: Option<String>,
    temperature_c: Option<i32>,
    power_on_hours: Option<u64>,
    reallocated_sectors: Option<u64>,
    wear_level_percent: Option<u8>,
}

fn read_smart(device: &str) -> Option<SmartInfo> {
    let mut cmd = Command::new("smartctl");
    cmd.arg("-a");
    if device.contains("nvme") {
        cmd.arg("-d").arg("nvme");
    }
    cmd.arg(device);
    let output = cmd.output().ok()?;
    let text = String::from_utf8_lossy(&output.stdout);
    if text.trim().is_empty() {
        return None;
    }

    let model = text
        .lines()
        .find(|l| l.starts_with("Device Model") || l.starts_with("Model Number"))
        .and_then(|l| l.split(':').nth(1))
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty());

    let temperature_c = text
        .lines()
        .find(|l| l.contains("Temperature_Celsius") || l.contains("Airflow_Temperature_Cel"))
        .and_then(|l| l.split_whitespace().last())
        .and_then(|v| v.parse().ok())
        .or_else(|| {
            text.lines()
                .find(|l| l.trim_start().starts_with("Temperature:"))
                .and_then(|l| l.split(':').nth(1))
                .and_then(|v| v.trim().split_whitespace().next())
                .and_then(|v| v.parse().ok())
        });

    let power_on_hours = text
        .lines()
        .find(|l| l.contains("Power_On_Hours"))
        .and_then(|l| l.split_whitespace().last())
        .and_then(|v| v.parse().ok())
        .or_else(|| {
            text.lines()
                .find(|l| l.trim_start().starts_with("Power On Hours:"))
                .and_then(|l| l.split(':').nth(1))
                .map(|v| v.trim().replace(',', ""))
                .and_then(|v| v.parse().ok())
        });

    let reallocated_sectors = text
        .lines()
        .find(|l| l.contains("Reallocated_Sector_Ct"))
        .and_then(|l| l.split_whitespace().last())
        .and_then(|v| v.parse().ok());
    let wear_level_percent = text
        .lines()
        .find(|l| l.contains("Wear_Leveling_Count"))
        .and_then(|l| l.split_whitespace().last())
        .and_then(|v| v.parse::<u8>().ok())
        .or_else(|| {
            text.lines()
                .find(|l| l.trim_start().starts_with("Percentage Used:"))
                .and_then(|l| l.split(':').nth(1))
                .map(|v| v.trim().trim_end_matches('%').to_string())
                .and_then(|v| v.parse().ok())
        });

    Some(SmartInfo { model, temperature_c, power_on_hours, reallocated_sectors, wear_level_percent })
}

fn evaluate(h: &DiskHealth) -> HealthStatus {
    if let Some(sectors) = h.reallocated_sectors {
        if sectors > 100 {
            return HealthStatus::Critical;
        }
        if sectors > 0 {
            return HealthStatus::Warning;
        }
    }
    if let Some(temp) = h.temperature_c {
        if temp > 65 {
            return HealthStatus::Critical;
        }
        if temp > 55 {
            return HealthStatus::Warning;
        }
    }
    if let Some(wear) = h.wear_level_percent {
        if wear > 90 {
            return HealthStatus::Critical;
        }
        if wear > 75 {
            return HealthStatus::Warning;
        }
    }

    let usage_ratio = 1.0 - (h.available_bytes as f64 / h.total_bytes.max(1) as f64);
    if usage_ratio > 0.97 {
        return HealthStatus::Critical;
    }
    if usage_ratio > 0.90 {
        return HealthStatus::Warning;
    }
    HealthStatus::Good
}