diff options
| -rw-r--r-- | src/cli.rs | 22 | ||||
| -rw-r--r-- | src/main.rs | 54 | ||||
| -rw-r--r-- | src/registers.rs | 1 | ||||
| -rw-r--r-- | src/trace.rs | 104 | ||||
| -rw-r--r-- | src/ui.rs | 131 |
5 files changed, 201 insertions, 111 deletions
diff --git a/src/cli.rs b/src/cli.rs new file mode 100644 index 0000000..685b26f --- /dev/null +++ b/src/cli.rs @@ -0,0 +1,22 @@ +use clap::Parser; + +#[derive(Parser)] +#[command(author, version, about, long_about = None)] +pub struct Args { + /// Command to execute from ptrace + #[arg(short, long)] + pub command: Option<String>, + + /// Attach the tracing to an existing process ID. We're using the `-p` short flag because + /// strace uses it + #[arg(short = 'p', long)] + pub attach: Option<i32>, + + /// Write the output to a file instead of the standard output + #[arg(short = 'f', long = "file")] + pub file_to_print: Option<String>, + + /// If defined, it hides the TUI + #[arg(long = "no-tui", default_value_t = false)] + pub no_tui: bool, +} diff --git a/src/main.rs b/src/main.rs index 19a9533..136c8c6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,50 +1,46 @@ +mod cli; mod registers; mod trace; mod ui; -use std::{ - io::{self, Write}, - str, -}; +use crate::cli::Args; use crate::trace::{exec, trace}; -use crate::ui::run_tui; +use crate::ui::UI; + use clap::Parser; use fork::{fork, Fork}; use nix::unistd::Pid; - -#[derive(Parser)] -#[command(author, version, about, long_about = None)] -struct Args { - /// Command to execute from ptrace - command: String, - - /// Write the output to a file instead of the standard output - #[arg(short = 'f', long = "file")] - file_to_print: Option<String>, - - /// If defined, it hides the TUI - #[arg(long = "no-tui", default_value_t = false)] - no_tui: bool, -} +use trace::attach; /// Create a fork of the program and execute the process in the child. Parent gets the pid /// value and trace it. fn main() -> anyhow::Result<()> { let args = Args::parse(); - let pid = match fork() { - Ok(Fork::Child) => return exec(&args.command), - Ok(Fork::Parent(child)) => Pid::from_raw(child), - Err(err) => panic!("fork() failed: {err}"), + let pid = if args.command.is_some() { + match fork() { + Ok(Fork::Child) => return exec(&args.command.unwrap()), + Ok(Fork::Parent(child)) => Pid::from_raw(child), + Err(err) => panic!("fork() failed: {err}"), + } + } else if args.attach.is_some() { + let pid = Pid::from_raw(args.attach.unwrap()); + + if attach(pid).is_ok() { + pid + } else { + panic!("Unable to attach to process {pid}"); + } + } else { + panic!("You must define a command or a PID to attach"); }; - let registers = trace(pid, args.file_to_print)?; if !args.no_tui { - run_tui(pid, ®isters)?; + let mut ui = UI::new(); + + ui.start(pid, &args)?; } else { - for line in registers { - writeln!(io::stdout(), "{}", line.output())?; - } + trace(pid, &args, true)?; } Ok(()) diff --git a/src/registers.rs b/src/registers.rs index d73424d..cf2aa98 100644 --- a/src/registers.rs +++ b/src/registers.rs @@ -6,6 +6,7 @@ use ratatui::{ }; /// Struct used to manipulate registers data from https://docs.rs/libc/0.2.147/libc/struct.user_regs_struct.html +#[derive(Debug)] pub struct RegistersData { orig_rax: u64, rdi: u64, diff --git a/src/trace.rs b/src/trace.rs index 3612519..c097fba 100644 --- a/src/trace.rs +++ b/src/trace.rs @@ -1,3 +1,4 @@ +use crate::cli::Args; use crate::registers::RegistersData; use nix::{ sys::{ @@ -7,7 +8,14 @@ use nix::{ }, unistd::Pid, }; -use std::{fs::File, io::Write, os::unix::process::CommandExt, process::Command}; +use std::{ + fs::File, + // fs::File, + io::{self, Write}, + os::unix::process::CommandExt, + process::Command, + str, +}; /// Exec the `command` value tracing it with `ptrace` lib pub fn exec(command: &str) -> anyhow::Result<()> { @@ -25,54 +33,82 @@ pub fn exec(command: &str) -> anyhow::Result<()> { Ok(()) } -/// Trace a process with `pid` ID and returns a list of `RegistersData` -pub fn trace(pid: Pid, file_to_print: Option<String>) -> anyhow::Result<Vec<RegistersData>> { - // Since you have to do 2 syscalls (start and end) you have to alternate the print value, - // because it could be equals except for the `rax` register. - let mut have_to_print = true; +/// Attach a ptrace status to a `pid` +pub fn attach(pid: Pid) -> anyhow::Result<()> { + ptrace::attach(pid)?; + + Ok(()) +} +/// Trace a process with `pid` ID and returns a list of `RegistersData` +pub fn trace(pid: Pid, args: &Args, run_loop: bool) -> anyhow::Result<Vec<RegistersData>> { // First wait for the parent process _ = waitpid(pid, None)?; - // If `fiole_to_print` is not None, create a new file with that value for redirecting all the + // FIXME: file writing on attachment + // If `file_to_print` is not None, create a new file with that value for redirecting all the // output (also in stdout) let mut f = None; - if let Some(filename) = file_to_print { + if let Some(filename) = &args.file_to_print { f = Some(File::create(filename)?); } let mut lines: Vec<RegistersData> = Vec::new(); - loop { - have_to_print ^= true; - ptrace::syscall(pid, None)?; - let status = waitpid(pid, None)?; + if run_loop { + // Since you have to do 2 syscalls (start and end) you have to alternate the print value, + // because it could be equals except for the `rax` register. + let mut have_to_print = true; - match status { - // Break the loop if the process exists - WaitStatus::Exited(_pid, _) => { - break; - } - // Match the stopped value for a process - WaitStatus::Stopped(pid, signal) => { - match signal { - Signal::SIGTRAP => { - if have_to_print { - let reg = RegistersData::new(ptrace::getregs(pid)?); - - if let Some(ref mut f) = f { - writeln!(f, "{}", reg.output())?; - } - - lines.push(reg); + loop { + match trace_next(pid)? { + Some(reg) => { + have_to_print ^= true; + if have_to_print { + if let Some(ref mut f) = f { + writeln!(f, "{}", reg.output())?; + } + + if args.no_tui { + writeln!(io::stdout(), "{}", reg.output())?; } + + lines.push(reg); } - _ => {} - }; + } + None => { + break; + } } - _ => {} - }; + } } - Ok(lines) } + +/// Get the next step for a ptrace process +pub fn trace_next(pid: Pid) -> anyhow::Result<Option<RegistersData>> { + ptrace::syscall(pid, None)?; + let status = waitpid(pid, None).unwrap(); + + match status { + // Match the stopped value for a process + WaitStatus::Stopped(pid, signal) => { + match signal { + Signal::SIGTRAP => { + let reg = RegistersData::new(ptrace::getregs(pid)?); + return Ok(Some(reg)); + } + _ => {} + }; + } + _ => {} + }; + + Ok(None) +} + +/// Kill a process traced by ptrace +pub fn trace_kill(pid: Pid) -> anyhow::Result<()> { + let _ = ptrace::kill(pid); + Ok(()) +} @@ -1,4 +1,8 @@ -use crate::registers::RegistersData; +use crate::{ + cli::Args, + registers::RegistersData, + trace::{trace, trace_kill, trace_next}, +}; use crossterm::{ event::{self, Event, KeyCode}, terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen}, @@ -8,19 +12,95 @@ use nix::unistd::Pid; use ratatui::{prelude::*, widgets::*}; use std::io::{self, stdout}; -struct UI { +pub struct UI { height: usize, max_lines: usize, scroll: usize, + lines: Vec<RegistersData>, } impl UI { - fn new() -> UI { + pub fn new() -> UI { UI { height: 0, max_lines: 0, scroll: 0, + lines: vec![], + } + } + + pub fn add_line(&mut self, registers: RegistersData) { + self.lines.push(registers); + self.max_lines = self.lines.len() + 1; + } + + pub fn get_paragraph(&self, pid: Pid) -> Paragraph { + let lines: Vec<Line> = self.lines.iter().map(|x| x.output_ui()).collect(); + let paragraph = Paragraph::new(lines) + .block( + Block::default() + .border_style(Style::default().fg(Color::Yellow)) + .title(format!("[{pid}]")) + .title( + block::Title::from(format!( + "[lines {}-{}]", + self.scroll, + self.scroll + self.height + )) + .position(block::Position::Bottom) + .alignment(Alignment::Right), + ) + .borders(Borders::ALL), + ) + .scroll((self.scroll as u16, 0)); + + paragraph + } + + pub fn start(&mut self, pid: Pid, args: &Args) -> anyhow::Result<()> { + enable_raw_mode()?; + stdout().execute(EnterAlternateScreen)?; + let mut terminal = Terminal::new(CrosstermBackend::new(stdout()))?; + + let mut have_to_print = true; + let mut have_to_trace = !args.command.is_some(); + let mut should_quit = false; + + let lines = trace(pid, &args, args.command.is_some())?; + if lines.len() > 1 { + for line in lines { + self.add_line(line); + } + } + while !should_quit { + if have_to_trace { + if let Some(reg) = trace_next(pid)? { + have_to_print ^= true; + if have_to_print { + self.add_line(reg); + } + } else { + have_to_trace = false; + } + } + + self.height = terminal.get_frame().size().height as usize; + terminal.draw(|frame| { + let size = frame.size(); + + frame.render_widget(self.get_paragraph(pid), size); + })?; + + should_quit = handle_events(self)?; } + + // FIXME: avoid this kill without Rust errors + let _ = trace_kill(pid); + + disable_raw_mode()?; + stdout().execute(LeaveAlternateScreen)?; + + Ok(()) } } @@ -55,48 +135,3 @@ fn handle_events(ui: &mut UI) -> io::Result<bool> { } Ok(false) } - -pub fn run_tui(pid: Pid, registers: &Vec<RegistersData>) -> anyhow::Result<()> { - enable_raw_mode()?; - stdout().execute(EnterAlternateScreen)?; - let mut terminal = Terminal::new(CrosstermBackend::new(stdout()))?; - - let mut ui = UI::new(); - ui.max_lines = registers.len() + 1; - - let mut should_quit = false; - while !should_quit { - ui.height = terminal.get_frame().size().height as usize; - terminal.draw(move |frame| { - let size = frame.size(); - let lines: Vec<Line> = registers.iter().map(|x| x.output_ui()).collect(); - - frame.render_widget( - Paragraph::new(lines) - .block( - Block::default() - .border_style(Style::default().fg(Color::Yellow)) - .title(format!("[{pid}]")) - .title( - block::Title::from(format!( - "[lines {}-{}]", - ui.scroll, - ui.scroll + ui.height - )) - .position(block::Position::Bottom) - .alignment(Alignment::Right), - ) - .borders(Borders::ALL), - ) - .scroll((ui.scroll as u16, 0)), - size, - ); - })?; - should_quit = handle_events(&mut ui)?; - } - - disable_raw_mode()?; - stdout().execute(LeaveAlternateScreen)?; - - Ok(()) -} |