nyanit/src/main.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/>.

use ratatui::{Frame, Terminal, backend::TermionBackend};
use ratatui::{
    layout::{Constraint, Layout},
    style::{Color, Style, Stylize},
    text::Line,
    widgets::{Block, Borders, List, ListState, Paragraph},
};

use ratatui_textarea::TextArea;

use std::io::{self, Read, stdout};
use std::process::ExitCode;
use std::thread;
use std::time::Duration;

use termion::event::{self, Event};
use termion::raw::IntoRawMode;

mod nyan_system;

use nyan_system::ExecCommandError;

const BEE_CAT: &str = include_str!("bee_cat.txt");

//const ROOT_PARTLABEL: &str = "bee-root";
const ROOT_PARTLABEL: &str = "bee-root";
const TMP_ROOT: &str = "/mnt/bee_root";
const TMP_ROOT_SNAPS: &str = "/mnt/bee-root/bzz_snaps";
const SYS_INIT: &str = "/sbin/init";

// fn bee_root_is_present() -> bool {
//    nyan_system::list_gpt_labeled_devs().iter().filter(|s| s.ends_with(ROOT_PARTLABEL)).count() > 0
// }

// fn main() -> ExitCode {
//     println!("Looking for {}", ROOT_PARTLABEL);
//     if bee_root_is_present() {
//         println!("{} found!", ROOT_PARTLABEL);
//     }
//     else {
//         print!("Cannot find {}!", ROOT_PARTLABEL);
//         print!("Exiting...");
//         return ExitCode::FAILURE;
//     }
//     // 1 get labels and see if there is bzzpss-root in them

//     ExitCode::SUCCESS
// }

#[derive(Debug, PartialEq)]
enum NyanitState {
    Init,
    WaitForF2(Vec<String>),
    ChooseSnapshot(Vec<String>),
    EnterPassword,
    Error(String),
    HandOff(String),
}

/// Mounting bee-root and chiking if there any snapshots
fn mount_and_check_snaps(bee_root_dev: &str) -> Result<NyanitState, ExecCommandError> {
    // mount bee-root
    nyan_system::mount_bee_root(TMP_ROOT, bee_root_dev)?;

    // list snaps if any
    let mut snaps_list = nyan_system::list_bzzpss_snaps(TMP_ROOT_SNAPS)?;
    // if any snaps detected whait for user to press F2
    if snaps_list.len() > 0 {
        snaps_list.insert(0, TMP_ROOT.to_string());
        return Ok(NyanitState::WaitForF2(snaps_list));
    }

    // else tell to finish boot
    Ok(NyanitState::HandOff(TMP_ROOT.into()))
}

fn init_stage() -> Result<(String, NyanitState), ExecCommandError> {
    // mount sysfss
    nyan_system::prepare_env().map_err(|ex| format!("prepare_env error: {}", ex))?;

    // get bee-root device name aka /dev/sda1
    let bee_root_dev = nyan_system::find_partlabel(ROOT_PARTLABEL)?;
    // check if device is encrypted
    if nyan_system::bzzpss_dev_encrypted(&bee_root_dev).map_err(|ex| format!("Error checking if device is encrypted: {}", ex))? {
        // if encrypted return EnterPassword so used get promted for pass
        return Ok((bee_root_dev, NyanitState::EnterPassword));
    }

    // if unencrypted
    let check_res = mount_and_check_snaps(&bee_root_dev)?;
    Ok((bee_root_dev, check_res))
}

struct NyanitTUI<'a> {
    exit: bool,
    state: NyanitState,
    bee_root_dev: String,
    list_state: ListState,
    time_counter: u32,
    textarea: Option<TextArea<'a>>,
}

impl NyanitTUI<'_> {
    pub fn new() -> io::Result<Self> {
        Ok(Self {
            exit: false,
            state: NyanitState::Init,
            bee_root_dev: String::new(),
            list_state: ListState::default().with_selected(Some(0)),
            time_counter: 0,
            textarea: None,
        })
    }

    /// runs the application's main loop until the user quits
    pub fn run(&mut self) -> io::Result<()> {
        self.state = match init_stage() {
            Ok(res) => {
                self.bee_root_dev = res.0;
                res.1
            },
            Err(ex) => {
                NyanitState::Error(format!("Error while trying to bzzz-pss things:\n{}", ex))
            },
        };

        // if drive unencrypted and there is no root snapshots finish boot
        match &self.state {
            // if drive unencrypted and there is no root snapshots finish boot
            NyanitState::HandOff(root_mnt) => {
                if let Err(ex) = nyan_system::hand_off_control(root_mnt, SYS_INIT) {
                    self.state = NyanitState::Error(ex.to_string())
                }
            }
            // init password input
            NyanitState::EnterPassword => {
                let mut textarea: TextArea<'_> = TextArea::default();
                textarea.set_cursor_line_style(Style::default());
                textarea.set_mask_char('\u{1F41D}'); //U+2022 BULLET (•)
                textarea.set_placeholder_text("Please enter your password");
                textarea.set_style(Style::default().fg(Color::LightYellow));
                textarea.set_block(Block::default().borders(Borders::ALL).title("Password"));
                self.textarea = Some(textarea);
            }
            _ => (),
        }

        // Init terminal
        // 1. Enable raw mode
        let stdout = stdout().into_raw_mode()?;
        // 2. Setup Terminal
        let backend = TermionBackend::new(stdout);
        let mut terminal = Terminal::new(backend)?;

        // init tty async input and start actually drawing things now!
        let mut tty_input_async = termion::async_stdin().bytes();

        self.time_counter = 0;
        while !self.exit {
            terminal.draw(|frame| self.draw(frame))?;
            self.handle_events(&mut tty_input_async)?;
        }
        Ok(())
    }

    fn draw(&mut self, frame: &mut Frame) {
        let area = frame.area();
        // 1. Create the outer block with borders
        let outer_block = Block::bordered()
            .title(Line::from(" BZZ-PSS ").centered())
            .light_yellow();
        // 2. Calculate the inner area excluding the borders
        let inner_area = outer_block.inner(area);
        // 3. Render the border block itself
        frame.render_widget(outer_block, area);

        // 1. Create vertical layout: Top (flexible), Center (fixed height), Bottom (flexible)
        let vertical_chunks = Layout::vertical([
            Constraint::Length(7),      // logo top part
            Constraint::Percentage(80), // Height of the paragraph block
            Constraint::Length(3),      // footer
        ])
        .split(inner_area);

        // 2. Create horizontal layout inside the center chunk: Left (flexible), Center (fixed width), Right (flexible)
        let horizontal_chunks = Layout::horizontal([
            Constraint::Fill(2),
            Constraint::Length(11),
            Constraint::Fill(2),
        ])
        .split(vertical_chunks[0]);

        // 3. Render the logo in the upper chunk
        let paragraph = Paragraph::new(BEE_CAT).left_aligned().light_yellow();
        frame.render_widget(paragraph, horizontal_chunks[1]);

        // now check app state
        match &self.state {
            NyanitState::Error(ex) => {
                // render big red wolf in the center with err messaghe
                let err_paragraph = Paragraph::new(ex.as_str()).centered().light_red().bold();
                frame.render_widget(err_paragraph, vertical_chunks[1]);
            }
            NyanitState::WaitForF2(_) => {
                // show press F2 for options text
                let f2_paragraph = Paragraph::new("Press F2 for options")
                    .centered()
                    .magenta()
                    .bold();
                frame.render_widget(f2_paragraph, vertical_chunks[1]);
            }
            NyanitState::ChooseSnapshot(snaps_list) => {
                // render menu for snap choosing
                frame.render_stateful_widget(
                    List::new(snaps_list.iter().map(String::as_str))
                        .block(Block::default().borders(Borders::ALL))
                        .highlight_symbol(">>")
                        .highlight_style(Style::new().bold()),
                    vertical_chunks[1],
                    &mut self.list_state,
                );
            }
            NyanitState::EnterPassword => {
                frame.render_widget(self.textarea.as_ref().unwrap(), vertical_chunks[1]);
            }
            NyanitState::HandOff(_) => {
                let ok_text = Paragraph::new("Booting\nKira Linux\n \u{1F41D} \u{1F41D} :3")
                    .block(Block::bordered().light_green())
                    .light_green();
                frame.render_widget(ok_text, vertical_chunks[1]);
                // just some time to look at message
                thread::sleep(Duration::from_millis(1000))
            }
            _ => (),
        };
    }

    fn handle_events(&mut self, tty_input: &mut io::Bytes<termion::AsyncReader>) -> io::Result<()> {
        // cheking if there is any input to read
        if let Some(item) = tty_input.next() {
            let read_byte = item?;

            let e = event::parse_event(read_byte, tty_input)?;
            match &self.state {
                // finishing boot
                NyanitState::HandOff(root_mnt) => {
                    if let Err(ex) = nyan_system::hand_off_control(root_mnt, SYS_INIT) {
                        self.state = NyanitState::Error(ex.to_string())
                    }
                }
                NyanitState::WaitForF2(vol_list) => {
                    if e == Event::Key(event::Key::F(2)) {
                        self.state = NyanitState::ChooseSnapshot(vol_list.clone());
                    }
                }
                NyanitState::ChooseSnapshot(vol_list) => match e {
                    Event::Key(event::Key::Down) => self.list_state.select_next(),
                    Event::Key(event::Key::Up) => self.list_state.select_previous(),
                    Event::Key(event::Key::Char('\n')) => {
                        if let Some(idx) = self.list_state.selected() {
                            self.state = NyanitState::HandOff(vol_list[idx].clone());
                        }
                    }
                    _ => (),
                },
                NyanitState::EnterPassword => {
                    match e {
                        // try to unlock if enter pressed
                        Event::Key(event::Key::Char('\n')) => {
                            let password = self.textarea.as_ref().unwrap().lines()[0].clone();
                            match nyan_system::bcachefs_unlock(password, &self.bee_root_dev) {
                                Ok(unlocked) => {
                                    if unlocked {
                                        // need to check for subvolumes and stuff
                                        self.state = match mount_and_check_snaps(&self.bee_root_dev) {
                                            Ok(st) => {
                                                self.time_counter = 0;
                                                st
                                            }
                                            Err(ex) => NyanitState::Error(ex.to_string()),
                                        }
                                    }
                                    // ask to enter password again
                                    else {
                                        self.textarea
                                            .as_mut()
                                            .unwrap()
                                            .set_placeholder_text("Wrong password, try again");
                                        self.textarea
                                            .as_mut()
                                            .unwrap()
                                            .set_style(Style::default().fg(Color::LightRed));
                                        self.textarea.as_mut().unwrap().clear();
                                    }
                                }
                                Err(ex) => {
                                    self.state = NyanitState::Error(ex.to_string());
                                }
                            }
                        }
                        _ => {
                            self.textarea
                                .as_mut()
                                .unwrap()
                                .set_style(Style::default().fg(Color::LightYellow));
                            self.textarea.as_mut().unwrap().input(e);
                        }
                    }
                }
                _ => (),
            }
        } else {
            // Tick time
            thread::sleep(Duration::from_millis(500));

            if let NyanitState::WaitForF2(_) = self.state {
                self.time_counter = self.time_counter.wrapping_add(1u32);
                // 2.5 sec wait
                if self.time_counter > 5 {
                    // boot from root
                    self.state = NyanitState::HandOff(TMP_ROOT.into());
                }
            }
        }
        Ok(())
    }
}



fn main() -> ExitCode {
    // nyan_system::set_env_vars_unsafe();
    let mut app = NyanitTUI::new().unwrap();
    app.run().unwrap();
    // Terminal restores automatically when `stdout` (RawTerminal) is dropped

    ExitCode::SUCCESS
}