1 files changed, 25 insertions, 29 deletions

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, &registers)?;
+        let mut ui = UI::new();
+
+        ui.start(pid, &args)?;
     } else {
-        for line in registers {
-            writeln!(io::stdout(), "{}", line.output())?;
-        }
+        trace(pid, &args, true)?;
     }
 
     Ok(())