nyanit/src/nyan_system.rs

// <Nyanit - init program fot initrd.>
// Copyright (C) <2026>  <Kira Foundation>

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

/*
   Functions fot system interaction
*/

use nix;
use std::fs;
use std::process::Command;

// timeout for block deices population
const WAIT_FOR_BLK_DEVS: u64 = 5u64;

#[derive(Debug)]
pub struct CommandFailed {
    pub stderr: Vec<u8>,
    pub cmd: String,
}

impl std::fmt::Display for CommandFailed {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(
            f,
            "Command failed! Command: {}; Stderr: {}",
            self.cmd,
            String::from_utf8(self.stderr.clone())
                .unwrap_or_else(|_| { "STDEER IS NOT VALID UTF-8".to_string() })
        )
    }
}

impl CommandFailed {
    pub fn new(cmd: &str, args: &Vec<&str>, stderr: Vec<u8>) -> Self {
        Self {
            stderr: stderr,
            cmd: format!("{} {}", cmd, args.join(" ")),
        }
    }
}

#[derive(Debug)]
pub enum ExecCommandError {
    IoError(std::io::Error),
    FromUtf8Error(std::string::FromUtf8Error),
    CommandFailed(CommandFailed),
    Other(String),
}

// 2. Implement From for each wrapped error type
impl From<std::io::Error> for ExecCommandError {
    fn from(err: std::io::Error) -> Self {
        ExecCommandError::IoError(err)
    }
}

impl From<std::string::FromUtf8Error> for ExecCommandError {
    fn from(err: std::string::FromUtf8Error) -> Self {
        ExecCommandError::FromUtf8Error(err)
    }
}

impl From<CommandFailed> for ExecCommandError {
    fn from(err: CommandFailed) -> Self {
        ExecCommandError::CommandFailed(err)
    }
}

impl From<String> for ExecCommandError {
    fn from(err: String) -> Self {
        ExecCommandError::Other(err)
    }
}

impl std::fmt::Display for ExecCommandError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            ExecCommandError::IoError(e) => write!(f, "IO Error: {}", e),
            ExecCommandError::FromUtf8Error(e) => {
                write!(f, "IO streem to String parse Error: {}", e)
            }
            ExecCommandError::CommandFailed(ex) => {
                write!(f, "{}", ex)
            }
            ExecCommandError::Other(msg) => write!(f, "Error: {}", msg),
        }
    }
}

// Implement std::error::Error
impl std::error::Error for ExecCommandError {}

/// exec given command with arguments, capturing stderr in case or failure
pub fn exec_command(cmd: &str, args: &Vec<&str>) -> Result<(), ExecCommandError> {
    let out = Command::new(cmd).args(args).output()?;
    if out.status.success() {
        Ok(())
    } else {
        Err(CommandFailed::new(cmd, args, out.stderr).into())
    }
}

pub fn try_find_partlabel(partlabel: &str, timeout_secs: u64) -> Option<String> {
    use std::{thread, time::Duration};

    let start = std::time::Instant::now();
    while start.elapsed().as_secs() < timeout_secs {
        if let Ok(Some(label)) = find_partlabel(partlabel) {
            return Some(label);
        }
        thread::sleep(Duration::from_millis(100));
    }
    None
}

/// Gets device node aka /dev/sda1 from gpr partition label
/// blkid -s PARTLABEL --match-token PARTLABEL="bzz-primary"
pub fn find_partlabel(partlabel: &str) -> Result<Option<String>, ExecCommandError> {
    let out = Command::new("/usr/bin/blkid")
        .args([
            "-s",
            "PARTLABEL",
            "-t",
            format!("PARTLABEL=\"{}\"", partlabel).as_str(),
        ])
        .output()?;
    if out.status.success() {
        let out_text = String::from_utf8(out.stdout)?;

        let split_res = out_text.split_once(':').ok_or_else(|| {
            ExecCommandError::Other(format!("Unexpected blkid output format: \"{}\"", out_text))
        })?;

        Ok(Some(split_res.0.to_string()))
    } else {
        if let Some(code) = out.status.code()
            && code == 2
        {
            Ok(None)
        } else {
            Err(ExecCommandError::Other(format!(
                "Error while looking for partlabel={}: \"{}\"",
                partlabel,
                String::from_utf8(out.stderr).unwrap_or_else(|ex| ex.to_string())
            )))
        }
    }
}

// pub fn list_dir(dirname: &str) -> Result<String, ExecCommandError> {
//     let res: String = fs::read_dir(dirname)?
//         .filter_map(|entry| entry.ok())
//         .map(|entry| entry.file_name().to_string_lossy().to_string())
//         .collect::<Vec<String>>()
//         .join("\n");

//     Ok(res)
// }

/// exec given command with arguments, returning stdout parsed to UTF-8 string
/// capturing stderr in case or failure
// pub fn exec_command_out(cmd: &str, args: &Vec<&str>) -> Result<String, ExecCommandError> {
//     let out = Command::new(cmd).args(args).output()?;
//     if out.status.success() {
//         Ok(String::from_utf8(out.stdout)?)
//     } else {
//         Err(CommandFailed::new(cmd, args, out.stderr).into())
//     }
// }

/// bcachefs unlock --file=/tmp/bzzpsspass.txt /dev/loop0somedev
/// Ok so. It is stupid, because bcacheutils far from ideal.
/// But if we provide wrong password from file or stdin this stuff plays stupid and asks for right password.
/// Sooo.. In order to know if unlock was sucessfull we need to check if program exited with Ok status,
/// or it shits to STDERR still waiting for stdin
/// Returns true if unlock successful, and false otherwise
pub fn bcachefs_unlock(pass: String, dev: &str) -> Result<bool, ExecCommandError> {
    use std::io::Read;
    use std::io::Write;
    use std::process::Stdio;
    // "-k", "session" is important as it stores key in sessyin keyring what allows
    // system to mount filesystem after switch_root
    let mut child = Command::new("/usr/bin/bcachefs")
        .args(&vec!["unlock", "-k", "session", dev])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()?;

    let mut c_stderr = child
        .stderr
        .take()
        .ok_or_else(|| "Unable to get child process STDERR".to_string())?;

    // Write password to STDIN in a limited scope
    {
        let mut c_stdin = child
            .stdin
            .take()
            .ok_or_else(|| "Unable to get child process STDIN".to_string())?;
        // neccessery to use thread or we may be deadlocking
        c_stdin.write_all(pass.as_bytes())?;
        let bytes_written = c_stdin.write(b"\n")?;
        println!("bytes_written: {}", bytes_written);
        // stdin is dropped here automatically when the block ends
    }

    let mut buffer = [0u8; 10];
    // This blocks.
    // if unlok fails it reads 10 bytes (piece of err message)
    // if sucess it reads 0 bytes and exit
    let read_b = c_stderr.read(&mut buffer)?;
    if read_b > 0 {
        println!("STDERR error not empty, password incorrect! Try to kill a child.");
        child.kill()?;
    }

    // Safe to wait now; child will see EOF
    let exit_status = child.wait()?;
    // if exit status is SUCCESS - unloking is successful
    Ok(exit_status.success())
}

/// Cheking if device is encripted
pub fn bzzpss_dev_encrypted(dev_name: &str) -> Result<bool, ExecCommandError> {
    match exec_command("/usr/bin/bcachefs", &vec!["unlock", "--check", dev_name]) {
        Ok(_) => Ok(true), // if status OK - dev encrypted
        Err(ExecCommandError::CommandFailed(_)) => Ok(false), // if bcachefs returned error - unencrypted
        Err(ex) => Err(ex),                                   // othervice its error
    }
}

//const TMP_ROOT: &str = "/mnt/bee-root";
pub fn list_bzzpss_snaps(snaps_dir: &str) -> Result<Vec<String>, ExecCommandError> {
    use regex::Regex;
    use std::path::Path;

    // root snapshoot in a form of bzz-yyyy-mm-ddTHH-MM-SS
    let re = Regex::new(r"^bzz-[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}-[0-9]{2}-[0-9]{2}$").unwrap();

    let mut res: Vec<String> = Vec::new();
    // if snaps dir path do not exist return empty list
    if !Path::new(snaps_dir).exists() {
        return Ok(res);
    }

    for maybe_entry in fs::read_dir(snaps_dir)? {
        let entry = maybe_entry?;
        if entry.file_type()?.is_dir() {
            if let Some(dir_name) = entry.file_name().to_str()
                && re.is_match(dir_name)
            {
                res.push(dir_name.to_string());
            }
        }
    }
    Ok(res)
}

// pulls /sys/class/block
fn wait_for_block_devices(timeout_secs: u64) -> bool {
    use std::path::Path;
    use std::{thread, time::Duration};

    let start = std::time::Instant::now();
    while start.elapsed().as_secs() < timeout_secs {
        // Check if any block device exists (e.g., sda, nvme0n1, vda)
        // You can refine this to check for a specific device needed for root
        if Path::new("/sys/class/block")
            .read_dir()
            .ok()
            .and_then(|mut d| d.next().is_some().then_some(true))
            .unwrap_or(false)
        {
            // Optional: Verify specific device needed for root exists
            // if Path::new("/dev/nvme0n1").exists() { return true; }
            return true;
        }
        thread::sleep(Duration::from_millis(100));
    }
    false
}

// mounting process
//# Mount essential virtual filesystems
// mount -t proc none /proc
// mount -t sysfs none /sys
// modprobe bcachefs
pub fn prepare_env() -> Result<(), String> {
    use nix::mount::{MsFlags, mount};

    use std::path::Path;

    // Mount proc filesystem
    // mount -t proc none /proc
    let proc_mount_flags =
        MsFlags::MS_NOSUID | MsFlags::MS_NODEV | MsFlags::MS_NOEXEC | MsFlags::MS_RELATIME;
    mount(
        Some("proc"),
        Path::new("/proc"),
        Some("proc"),
        proc_mount_flags,
        Some("mode=0555"),
    )
    .map_err(|ex| format!("Error mounting /proc: {}", ex.to_string()))?;

    // mount -t sysfs none /sys
    let sys_mount_flags =
        MsFlags::MS_NOSUID | MsFlags::MS_NODEV | MsFlags::MS_NOEXEC | MsFlags::MS_RELATIME;
    mount(
        Some("sysfs"),
        Path::new("/sys"),
        Some("sysfs"),
        sys_mount_flags,
        Some("mode=0555"),
    )
    .map_err(|ex| format!("Error mounting /sys: {}", ex.to_string()))?;

    //waiting for block devs population
    if wait_for_block_devices(WAIT_FOR_BLK_DEVS) == false {
        return Err("Waiting for block devices to get populated timed out.".to_string());
    }

    let dev_mount_flags = MsFlags::MS_NOSUID | MsFlags::MS_RELATIME;
    mount(
        Some("devtmpfs"),
        "/dev",
        Some("devtmpfs"),
        dev_mount_flags,
        Some("mode=0755"),
    )
    .map_err(|ex| format!("Error mounting /dev: {}", ex.to_string()))?;

    // modprobe bcachefs
    // exec_command("modprobe", &vec!["bcachefs"])?;
    Ok(())
}

///
///  MUST be done before spawning any threads or async runtimes
// pub fn set_env_vars_unsafe() -> () {
//     use std::env;
//     // 1. Define your custom paths
//     let custom_bin = "/usr/bin";
//     let custom_lib = "/usr/lib";

//     // 2. Retrieve existing paths to append/prepend (optional but recommended)
//     // 3. Construct new path strings
//     let new_path = match env::var("PATH") {
//         Err(_) => custom_bin.to_string(),
//         Ok(v) => format!("{}:{}", custom_bin, v),
//     };

//     let new_ld_path = match env::var("LD_LIBRARY_PATH") {
//         Err(_) => custom_lib.to_string(),
//         Ok(v) => format!("{}:{}", custom_lib, v),
//     };

//     // 4. Set environment variables (UNSAFE in Rust 2024+)
//     // MUST be done before spawning any threads or async runtimes
//     unsafe {
//         env::set_var("PATH", new_path);
//         env::set_var("LD_LIBRARY_PATH", new_ld_path);
//     }
// }

// # Mount the real root filesystem
pub fn mount_bee_root(mount_point: &str, dev_name: &str) -> Result<(), ExecCommandError> {
    exec_command("/usr/bin/mount", &vec!["-o", "ro", dev_name, mount_point])?;
    Ok(())
}

// finish boot
pub fn hand_off_control(mnt_root: &str, init_path: &str) -> Result<(), ExecCommandError> {
    use std::env;
    use std::os::unix::process::CommandExt;

    // 2. Change directory to the new root (Mandatory for switch_root)
    env::set_current_dir(mnt_root)?;

    // # Clean up virtual filesystems
    // !!! do not do this, switch_root handle this stuff and fail if thay not present!
    // exec_command("umount", &vec!["/proc"])?;
    // exec_command("umount", &vec!["/sys"])?;

    // # Hand off control to the real system's init
    // exec switch_root /mnt/root /sbin/init
    // 3. Execute switch_root.
    // This call replaces the current Rust process.
    // It only returns if an error occurs.
    let err = Command::new("/usr/bin/switch_root")
        .arg(mnt_root)
        .arg(init_path)
        .exec();

    // If execution reaches here, exec() failed.
    Err(format!("switch_root failed: {}", err).into())
}