statusline 0.23.1

Simple and fast bash PS1 line with useful features
use crate::file;
use std::{
    fs::File,
    io::{BufRead as _, BufReader},
    path::Path,
};

#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
fn detect_vm_cpuid() -> Option<bool> {
    let cpuid = raw_cpuid::CpuId::new();
    let vendor_info = cpuid.get_vendor_info()?;
    match vendor_info.as_str() {
        // Known to only belong to hypervisors
        "XenVMMXenVMM" | "KVMKVMKVM\0\0\0" | "Linux KVM Hv" | "TCGTCGTCGTCG" | "VMwareVMware"
        | "Microsoft Hv" | "bhyve bhyve " | "QNXQVMBSQG\0\0" | "ACRNACRNACRN" | "SRESRESRESRE"
        | "MicrosoftXTA" | "VirtualApple" | "PowerVM Lx86" | "Neko Project" | "Insignia 586"
        | "Compaq FX!32" | "ConnectixCPU" => Some(true),

        // Known to only belong to hardware
        "GenuineIntel" | "AuthenticAMD" | "CentaurHauls" | "CyrixInstead" | "GenuineIotel"
        | "TransmetaCPU" | "GenuineTMx86" | "Geode by NSC" | "NexGenDriven" | "RiseRiseRise"
        | "SiS SiS SiS " | "UMC UMC UMC " | "Vortex86 SoC" | "  Shanghai  " | "HygonGenuine"
        | "Genuine  RDC" | "E2K MACHINE " | "VIA VIA VIA " | "AMD ISBETTER" => Some(false),

        other => {
            eprintln!("CPUID returned unknown vendor info: {other:?}");
            None
        }
    }
}

#[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))]
fn detect_vm_cpuid() -> Option<bool> {
    None
}

fn detect_vm_device_tree() -> Option<bool> {
    if Path::new("/proc/device-tree/hypervisor/compatible").exists() {
        return Some(true);
    }

    if Path::new("/proc/device-tree/ibm,partition-name").exists()
        && Path::new("/proc/device-tree/hmc-managed?").exists()
        && !Path::new("/proc/device-tree/chosen/qemu,graphic-width").exists()
    {
        return Some(true);
    }

    if let Ok(true) = file::exists_that("/proc/device-tree", |name| name.contains("fw-cfg")) {
        return Some(true);
    }

    if let Ok(s) = std::fs::read_to_string("/proc/device-tree/compatible")
        && s == "qemu,pseries"
    {
        return Some(true);
    }

    None
}

fn detect_vm_dmi_vendor() -> Option<bool> {
    [
        "/sys/class/dmi/id/product_name",
        "/sys/class/dmi/id/sys_vendor",
        "/sys/class/dmi/id/board_vendor",
        "/sys/class/dmi/id/bios_vendor",
        "/sys/class/dmi/id/product_version",
    ]
    .iter()
    .map(std::fs::read_to_string)
    .filter_map(Result::ok)
    .any(|name| {
        [
            "KVM",
            "OpenStack",
            "KubeVirt",
            "Amazon EC2",
            "QEMU",
            "VMware",
            "VMW",
            "innotek GmbH",
            "VirtualBox",
            "Xen",
            "Bochs",
            "Parallels",
            "BHYVE",
            "Hyper-V",
            "Apple Virtualization",
        ]
        .iter()
        .any(|vendor| name.starts_with(vendor))
    })
    .then_some(true)
}

fn detect_vm_smbios() -> Option<bool> {
    // See 7.1.2.2 "BIOS Characteristics Extension Byte 2" at [SMBIOS spec]
    // [SMBIOS spec]: https://www.dmtf.org/sites/default/files/standards/documents/DSP0134_3.4.0.pdf
    if let Ok(s) = std::fs::read("/sys/firmware/dmi/entries/0-0/raw")
        && let Some(byte) = s.get(0x13)
        && (byte >> 4_i32) & 1 == 1
    {
        return Some(true);
    }
    None
}

fn detect_vm_xen() -> Option<bool> {
    if !Path::new("/proc/xen").exists() {
        return None;
    }

    if let Ok(s) = std::fs::read_to_string("/sys/hypervisor/properties/features")
        && let Ok(val) = u64::from_str_radix(&s, 16)
    {
        return Some((val >> 11) & 1 != 1);
    }

    if let Ok(s) = std::fs::read_to_string("/proc/xen/capabilities")
        && s.contains("control_d")
    {
        return Some(false);
    }

    None
}

pub fn detect_vm() -> bool {
    if let Ok(file) = File::open("/proc/cpuinfo")
        && BufReader::new(file)
            .lines()
            .map_while(Result::ok)
            .any(|line| line.starts_with("vendor_id\t: User Mode Linux"))
    {
        return true;
    }

    if Path::new("/sys/hypervisor/type").exists() || Path::new("/proc/sysinfo").exists() {
        return true;
    }

    [
        detect_vm_cpuid,
        detect_vm_xen,
        detect_vm_smbios,
        detect_vm_dmi_vendor,
        detect_vm_device_tree,
    ]
    .iter()
    .find_map(|f| f())
    .unwrap_or(false)
}

pub fn detect_container() -> bool {
    if Path::new("/proc/vz").exists() && !Path::new("/proc/bc").exists()
        || Path::new("/run/host/container-daemon").exists()
        || Path::new("/run/systemd/container").exists()
        || Path::new("/run/.containerenv").exists()
        || Path::new("/.dockerenv").exists()
    {
        return true;
    }

    if let Ok(s) = std::fs::read_to_string("/proc/sys/kernel/osrelease")
        && (s.contains("Microsoft") || s.contains("WSL"))
    {
        return true;
    }

    if let Ok(file) = File::open("/proc/self/status")
        && let Some(pid) = BufReader::new(file)
            .lines()
            .map_while(Result::ok)
            .find_map(|line| line.strip_prefix("TracerPid:\t")?.parse::<usize>().ok())
        && let Ok(s) = std::fs::read_to_string(format!("/proc/{pid}/comm"))
        && s.starts_with("proot")
    {
        return true;
    }

    if let Ok(file) = File::open("/proc/1/environ")
        && BufReader::new(file)
            .split(0)
            .map_while(Result::ok)
            .any(|line| line.starts_with(b"container="))
    {
        return true;
    }

    false
}