meka 0.2.6

//! Process handling and monitoring using ptrace

use nix::libc::{self, SIGPIPE, SIGTRAP, sleep};
use nix::sys::signal::Signal;
use nix::unistd::{ForkResult, Pid, close, dup2_stderr, dup2_stdout, execve, fork, getpid};
use nix::{
    libc::SIGSTOP,
    sys::{
        memfd, ptrace,
        wait::{WaitPidFlag, WaitStatus, waitpid},
    },
};
use std::io;
use std::{
    collections::BTreeMap,
    env::{self, set_current_dir},
    ffi::{CString, NulError},
    path::PathBuf,
};

use crate::{
    args::RunArgs,
    errors::MekaError,
    meka::Meka,
    utils::{ParsedCommand, Proc, ProgramData, parse_cmd, read_memfd},
};

/// Process Options
#[derive(Debug, Clone)]
pub struct ProcessOptions {
    /// Capture environment
    pub capture_environ: bool,
    /// Capture status
    pub capture_status: bool,
}

/// Process Options
#[derive(Debug, Clone)]
pub struct RunOptions {
    /// Process cwd
    pub cwd: Option<PathBuf>,
    /// Wrap in a shell
    pub wrap_shell: bool,
    /// Write on insert
    pub write_on_insert: bool,
}

impl Meka {
    /// Process a command and return a map of process data
    /// # Errors
    /// Returns a `MekaError` if the command cannot be processed
    pub fn process_command(
        &self,
        path_var: Vec<PathBuf>,
        all_args: Vec<String>,
        run_options: RunArgs,
    ) -> Result<BTreeMap<i32, ProgramData>, MekaError> {
        let wrap_shell = run_options.wrap_shell;
        let process_options = ProcessOptions {
            capture_environ: run_options.capture_environ,
            capture_status: run_options.capture_status,
        };
        let ParsedCommand {
            envs: cmd_envs,
            executable: path_executable,
            args: args_cmd,
        } = parse_cmd(&all_args, &path_var, wrap_shell).map_err(|e| {
            if e.contains("not found") && self.should_show_debug(1) {
                let debug_infos = {
                    // not using env::join_paths()
                    let paths = path_var
                        .iter()
                        .map(|p| p.to_string_lossy())
                        .collect::<Vec<_>>()
                        .join(":");
                    format!(" (PATH is '{paths}')")
                };
                MekaError::new(format!("{e}{debug_infos}"))
            } else {
                MekaError::new(e)
            }
        })?;
        let path_binary_cstr = CString::new(path_executable.to_str().ok_or(MekaError::new(
            format!("Conversion of {} failed", path_executable.display()),
        ))?)
        .map_err(|e| {
            MekaError::new(format!(
                "Cannot create a CString from {}: {e}",
                path_executable.display()
            ))
        })?;

        let display = format!(
            "[{}]",
            args_cmd
                .iter()
                .map(|s| format!("\"{}\"", s)) // add quotes around each string
                .collect::<Vec<_>>()
                .join(", ")
        );
        self.logger1(format!(
            "Executing '{}' with args: {}",
            path_executable.display(),
            display
        ));

        let all_args_cstr = args_cmd
            .iter()
            .map(|s| CString::new(s.as_str()))
            .collect::<Result<Vec<CString>, NulError>>()
            .map_err(|e| {
                MekaError::new(format!(
                    "Cannot create CString(s) for args {}: {e}",
                    display
                ))
            })?;

        let mut process_table = BTreeMap::<i32, ProgramData>::new();

        let fd_out = memfd::memfd_create("out", memfd::MFdFlags::MFD_CLOEXEC)?;

        let fd_err = memfd::memfd_create("err", memfd::MFdFlags::MFD_CLOEXEC)?;
        if let Some(pwd) = &run_options.cwd {
            set_current_dir(pwd).map_err(|e| {
                MekaError::new(format!(
                    "Cannot set the current dir to {}: {e}",
                    pwd.display()
                ))
            })?;
        }
        // other option was to use unsafe set_var
        let env_vars: Vec<String> = env::vars()
            .map(|(key, val)| format!("{key}={val}"))
            .collect();
        let cmd_envs: Vec<String> = cmd_envs.iter().map(|(k, v)| format!("{k}={v}")).collect();
        let env_vars: Vec<String> = env_vars.into_iter().chain(cmd_envs).collect();
        let all_envs_cstr = env_vars
            .iter()
            .map(|s| CString::new(s.as_str()))
            .collect::<Result<Vec<CString>, NulError>>()
            .map_err(|e| MekaError::new(format!("Cannot create CString(s) for envs {e}")))?;
        match unsafe { fork() } {
            Ok(ForkResult::Child) => {
                // dup2_stdout(&fd_out)?;
                // dup2_stderr(&fd_err)?;
                // close(fd_out)?;
                // close(fd_err)?;
                ptrace::traceme()?;
                execve(&path_binary_cstr, &all_args_cstr, &all_envs_cstr)?;
                unreachable!();
            }
            Ok(ForkResult::Parent { child: child_pid }) => {
                let father_process_pid = getpid().as_raw();
                self.logger1(format!(
                    "Parent PID: {father_process_pid}, Child PID: {child_pid}"
                ));

                match waitpid(child_pid, None)? {
                    WaitStatus::Stopped(pid, sig) => {
                        self.logger2(format!(
                            "Child process {} stopped (signal={:?})",
                            pid.as_raw(),
                            sig
                        ));
                        // insert the parent process
                        let program = ProgramData::from_proc(
                            &Proc::new(pid.as_raw()),
                            process_options.clone(),
                        );
                        process_table.insert(pid.as_raw(), program);
                        if run_options.write_on_insert {
                            Self::write_meka_file(&process_table, &run_options.output_file)?;
                        }
                    }
                    WaitStatus::Exited(pid, return_code) => {
                        return Err(MekaError::new(format!(
                            "Failed to start the exec: pid({}) returned with code={} (it seems that {} is not an executable)",
                            pid.as_raw(),
                            return_code,
                            path_executable.display()
                        )));
                    }
                    status => {
                        return Err(MekaError::new(format!(
                            "Unexpected return from wait(): {:?}",
                            status
                        )));
                    }
                }

                fn set_ptrace_options(pid: libc::pid_t) -> io::Result<()> {
                    // combine options
                    let opts = libc::PTRACE_O_TRACESYSGOOD
                        | libc::PTRACE_O_TRACEFORK
                        | libc::PTRACE_O_TRACEVFORK
                        | libc::PTRACE_O_TRACECLONE
                        | libc::PTRACE_O_TRACEEXEC
                        | libc::PTRACE_O_TRACEEXIT;
                    let r =
                        unsafe { libc::ptrace(libc::PTRACE_SETOPTIONS, pid, 0, opts as c_long) };
                    if r != 0 {
                        return Err(io::Error::last_os_error());
                    }
                    eprintln!("[MON] set options for {}", pid);
                    Ok(())
                }

                fn resume_syscall(pid: libc::pid_t, sig_to_deliver: i32) -> io::Result<()> {
                    // use PTRACE_SYSCALL to continue to next syscall boundary, optionally delivering a signal
                    let r = unsafe {
                        libc::ptrace(libc::PTRACE_SYSCALL, pid, 0, sig_to_deliver as c_long)
                    };
                    if r != 0 {
                        return Err(io::Error::last_os_error());
                    }
                    Ok(())
                }
                use libc::{c_long, c_void};
                fn get_event_msg(pid: libc::pid_t) -> io::Result<u64> {
                    let mut msg: u64 = 0;
                    let r = unsafe {
                        libc::ptrace(
                            libc::PTRACE_GETEVENTMSG,
                            pid,
                            0,
                            &mut msg as *mut _ as *mut c_void,
                        )
                    };
                    if r != 0 {
                        return Err(io::Error::last_os_error());
                    }
                    Ok(msg)
                }

                // Set options for the initial pid immediately
                set_ptrace_options(child_pid.as_raw())
                    .map_err(|e| MekaError::new(format!("{e}")))?;
                resume_syscall(child_pid.as_raw(), 0)
                    .map_err(|e| MekaError::new(format!("{e}")))?;
                loop {
                    eprintln!("LLLOOO");
                    let mut raw_status: libc::c_int = 0;
                    let waited =
                        unsafe { libc::waitpid(-1, &mut raw_status as *mut _, libc::__WALL) };
                    eprintln!("LLLOOOPPPP");
                    if waited < 0 {
                        let e = io::Error::last_os_error();
                        // If EINTR, just continue waiting
                        if e.raw_os_error() == Some(libc::EINTR) {
                            continue;
                        }
                        return Err(MekaError::new(format!("{e}")));
                    }
                    let pid = waited as libc::pid_t;

                    // Exited normally
                    if unsafe { libc::WIFEXITED(raw_status) } {
                        let code = unsafe { libc::WEXITSTATUS(raw_status) };
                        eprintln!("[MON] pid {} exited with {}", pid, code);
                        if pid == child_pid.as_raw() {
                            break;
                        }
                        // If that was the original traced process and everything is done, we could break.
                        // For demo, keep monitoring other processes.
                        continue;
                    }

                    // Killed by signal
                    if unsafe { libc::WIFSIGNALED(raw_status) } {
                        let sig = unsafe { libc::WTERMSIG(raw_status) };
                        let core = unsafe { libc::WCOREDUMP(raw_status) };
                        eprintln!("[MON] pid {} killed by signal {} (core={})", pid, sig, core);
                        continue;
                    }

                    // Continued (WCONTINUED)
                    if unsafe { libc::WIFCONTINUED(raw_status) } {
                        eprintln!("[MON] pid {} continued", pid);
                        // resume without injecting a signal
                        resume_syscall(pid, 0).map_err(|e| MekaError::new(format!("{e}")))?;
                        continue;
                    }

                    // Stopped (including ptrace events)
                    if unsafe { libc::WIFSTOPPED(raw_status) } {
                        let wstopsig = unsafe { libc::WSTOPSIG(raw_status) }; // may include 0x80 from TRACESYSGOOD
                        eprintln!(
                            "[MON] pid {} stopped raw WSTOPSIG=0x{:x} ({})",
                            pid, wstopsig, wstopsig
                        );

                        // Check for ptrace event: some kernels put event number in the high bits of raw_status
                        let event = ((raw_status as u32) >> 16) as i32;
                        if event != 0 {
                            // This is a PTRACE_EVENT stop. Read event message (e.g., new child pid).
                            eprintln!("[MON] pid {} ptrace event {}", pid, event);
                            match get_event_msg(pid) {
                                Ok(msg) => {
                                    eprintln!("[MON] event message = {}", msg);
                                    // if fork/clone/vfork, event message typically contains child pid
                                    if event == libc::PTRACE_EVENT_FORK
                                        || event == libc::PTRACE_EVENT_VFORK
                                        || event == libc::PTRACE_EVENT_CLONE
                                    {
                                        let child = msg as libc::pid_t;
                                        if child > 0 {
                                            // set options for the child immediately
                                            if let Err(e) = set_ptrace_options(child) {
                                                eprintln!(
                                                    "[MON] failed to setoptions for child {}: {}",
                                                    child, e
                                                );
                                            }
                                        }
                                    }
                                }
                                Err(e) => {
                                    eprintln!("[MON] failed to get event msg for {}: {}", pid, e);
                                }
                            }
                            // After handling event, resume. If wstopsig includes 0x80 it's a syscall stop,
                            // otherwise treat as plain trap/signal (deliver none).
                            if (wstopsig & 0x80) != 0 {
                                // syscall-enter/exit stop
                                resume_syscall(pid, 0)
                                    .map_err(|e| MekaError::new(format!("{e}")))?;
                            } else if wstopsig == libc::SIGTRAP {
                                // plain SIGTRAP (internal)
                                resume_syscall(pid, 0)
                                    .map_err(|e| MekaError::new(format!("{e}")))?;
                            } else {
                                resume_syscall(pid, wstopsig)
                                    .map_err(|e| MekaError::new(format!("{e}")))?;
                            }
                            continue;
                        }

                        // Not an event: decide whether this is a syscall stop (TRACESYSGOOD -> 0x80 bit)
                        if (wstopsig & 0x80) != 0 {
                            // Syscall-enter or -exit stop
                            eprintln!("[MON] pid {} syscall stop (0x80 bit set)", pid);
                            resume_syscall(pid, 0).map_err(|e| MekaError::new(format!("{e}")))?;
                            continue;
                        }

                        // Plain SIGTRAP (no 0x80). Could be exec/bkpt; treat as internal.
                        if wstopsig == libc::SIGTRAP {
                            eprintln!("[MON] pid {} plain SIGTRAP", pid);
                            resume_syscall(pid, 0).map_err(|e| MekaError::new(format!("{e}")))?;
                            continue;
                        }

                        // Otherwise this is a real signal delivery (SIGPIPE, etc.). Forward it.
                        eprintln!("[MON] pid {} real signal delivery: {}", pid, wstopsig);

                        let r = unsafe {
                            libc::ptrace(libc::PTRACE_SINGLESTEP, pid, 0, wstopsig as c_long)
                        };
                        if r != 0 {
                            return Err(io::Error::last_os_error())
                                .map_err(|e| MekaError::new(format!("{e}")));
                        }
                        continue;
                    }

                    // Shouldn't reach here but be defensive
                    eprintln!(
                        "[MON] unexpected status for pid {} raw=0x{:x}",
                        pid, raw_status
                    );
                }
                self.logger(
                    1,
                    format!("[MONITOR] Child process has exited. PID: {child_pid}"),
                );
                if let Some(data) = read_memfd(fd_out) {
                    self.logger(2, "-------- Child stdout --------");
                    self.logger(2, &data);
                    self.logger(2, "-------- End --------");
                }
                if let Some(data) = read_memfd(fd_err) {
                    self.logger(2, "-------- Child stderr --------");
                    self.logger(2, &data);
                    self.logger(2, "-------- End --------");
                }
                self.logger(1, "[MONITOR] Parent process has finished.");
                Ok(process_table)
            }
            Err(_) => Err(MekaError::new("fork() failed")),
        }
    }
}