meka 0.2.6

//! Main meka struct and methods

use serde_json::json;

use crate::{
    args::{AnalyzeArgs, Commands, MekaArgs, RunArgs},
    errors::MekaError,
    meka_file::MekaFile,
    utils::{ProgramData, convert_to_relative, correct_args, get_path_var},
};
use std::{collections::BTreeMap, path::PathBuf};

/// Main struct for meka
#[derive(Default)]
pub struct Meka {
    /// Debug level (0 = no debug, higher = more debug)
    debug_level: u8,
}

/// Default meka output file
pub const MEKA_FILE: &str = "meka_recording.json";

impl Meka {
    /// Build meka from args
    pub fn from_args(args: &MekaArgs) -> Self {
        Meka::default().with_debug(args.debug)
    }

    /// Create a new Meka instance
    /// # Errors
    /// Returns `MekaError` if there is an error
    pub fn run_from_args(args: MekaArgs) -> Result<(), MekaError> {
        let meka = Self::from_args(&args);
        match args.command {
            Commands::Record(run_args) => meka.meka_run(run_args),
            Commands::Analyze(analysis_args) => meka.meka_analyze(analysis_args),
        }
    }

    /// Set the debug level (0 = no debug, higher = more debug)
    pub fn with_debug(mut self, level: u8) -> Self {
        self.debug_level = level;
        self
    }

    /// Should show debug
    #[inline]
    pub(crate) fn should_show_debug(&self, debug_level: u8) -> bool {
        debug_level <= self.debug_level
    }

    /// Log a message if the debug level is high enough
    #[inline]
    pub(crate) fn logger<P>(&self, debug_level: u8, msg: P)
    where
        P: AsRef<str>,
    {
        let msg = msg.as_ref();
        if self.should_show_debug(debug_level) {
            println!("[meka] {msg}");
        }
    }

    /// Process a command and return a map of process data
    /// # Errors
    /// Returns `MekaError` if there is an error
    pub fn meka_run(&self, run_args: RunArgs) -> Result<(), MekaError> {
        if run_args.write_on_insert && run_args.output_file == "-" {
            return Err(MekaError::new(
                "Cannot have both --write-on-insert and --output-file=-",
            ));
        }
        let print_json = run_args.json_output;
        let relative_pid = run_args.relative_pid;
        let output_file = run_args.output_file.clone();
        let path_var = run_args.path_var.clone().unwrap_or(get_path_var());
        let command = correct_args(&run_args.inputs);
        let map = match self.process_command(path_var, command, run_args) {
            Ok(map) => map,
            Err((e, map)) => {
                self.logger1(format!(
                    "Tracing failed {e} - trying to save the processes table"
                ));
                Self::write_meka_file(&map, &output_file)?;
                return Err(e);
            }
        };
        let map = if relative_pid {
            convert_to_relative(map)
        } else {
            map
        };
        let meka_file_value = json!({ "processes": map });
        if print_json || output_file == "-" {
            let json_data = serde_json::to_string_pretty(&meka_file_value)
                .map_err(|e| MekaError::new(format!("Cannot stringify json {e}")))?;
            println!("{json_data}");
        }
        if output_file != "-" {
            MekaFile::write_value(meka_file_value, &PathBuf::from(output_file))?;
        }

        Ok(())
    }

    /// Write meka file
    /// # Errors
    /// Fails if cannot write
    pub fn write_meka_file(
        map: &BTreeMap<i32, ProgramData>,
        output_file: &String,
    ) -> Result<(), MekaError> {
        let meka_file_value = json!({ "processes": map });
        MekaFile::write_value(meka_file_value, &PathBuf::from(output_file))?;
        Ok(())
    }

    /// Analyze a meka recording file and print a tree of processes
    /// # Errors
    /// Returns `MekaError` if there is an error
    pub fn meka_analyze(&self, analyze_args: AnalyzeArgs) -> Result<(), MekaError> {
        let input_file = analyze_args.input_file;
        let min_cmd_size = analyze_args.min_cmd_size;
        let meka_file = MekaFile::try_from_file(PathBuf::from(&input_file))?;
        let map = meka_file.processes;
        if map.is_empty() {
            return Err(MekaError::new("processes map is empty"));
        }
        // find index where parent_pid is 0
        let root = map
            .iter()
            .find(|(_, v)| v.parent_pid == 0)
            .ok_or(MekaError::new(format!(
                "Cannot find root process in {}",
                input_file
            )))?;

        let mut stack: Vec<(&i32, usize)> = vec![(root.0, 0)];

        while let Some((pid, space)) = stack.pop() {
            if let Some(proc) = map.get(pid) {
                let cmd = if proc.after_exec.is_empty() {
                    proc.cmdline.join(" ")
                } else {
                    // this is only true for fork-exec
                    // it can be cmdline then after_exec for multiples exec
                    proc.after_exec
                        .iter()
                        .map(|cmd| cmd.join(" "))
                        .collect::<Vec<String>>()
                        .join(" THEN ")
                };
                println!(
                    "{}{} - {}",
                    " ".repeat(space),
                    proc.pid,
                    &cmd.replace('\n', " ")[..min_cmd_size.min(cmd.len())],
                );

                for child_pid in proc.subprograms.iter().rev() {
                    stack.push((child_pid, space + 4));
                }
            }
        }

        // print tree
        Ok(())
    }
}

/// make_logger
macro_rules! make_logger {
    ($($name:ident => $level:expr),* $(,)?) => {
        $(
            #[inline]
            pub(crate) fn $name<P>(&self, msg: P)
            where
                P: AsRef<str>,
            {
                self.logger($level, msg);
            }
        )*
    }
}

impl Meka {
    make_logger! {
        logger1 => 1,
        logger2 => 2,
        logger3 => 3,
        logger4 => 4,
        // logger5 => 5,
        logger6 => 6,
    }
}