From 5c3c5033b36986641603253400e86850dc60c5e2 Mon Sep 17 00:00:00 2001
From: Santo Cariotti <santo@dcariotti.me>
Date: Sat, 28 Oct 2023 09:47:42 +0200
Subject: Avoid panics in main

---
 Cargo.lock  |  2 +-
 Cargo.toml  |  2 +-
 src/main.rs | 32 ++++++++++++++++++++------------
 3 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index bb222c8..afca903 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -434,7 +434,7 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
 
 [[package]]
 name = "sigma-trace"
-version = "0.0.1-c"
+version = "0.0.2"
 dependencies = [
  "anyhow",
  "byteorder",
diff --git a/Cargo.toml b/Cargo.toml
index 9b22e51..ae23863 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "sigma-trace"
-version = "0.0.1-c"
+version = "0.0.2"
 edition = "2021"
 authors = ["Santo Cariotti <santo@dcariotti.me>"]
 repository = "https://github.com/boozec/sigma"
diff --git a/src/main.rs b/src/main.rs
index f03c271..84ba250 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,37 +11,45 @@ use crate::ui::UI;
 use clap::Parser;
 use fork::{fork, Fork};
 use nix::unistd::Pid;
+use owo_colors::OwoColorize;
 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 = if args.command.is_some() {
+    let process: Result<Pid, String> = 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}"),
+            Ok(Fork::Parent(child)) => Ok(Pid::from_raw(child)),
+            Err(err) => Err(format!("fork() failed: {err}")),
         }
     } else if args.attach.is_some() {
         let pid = Pid::from_raw(args.attach.unwrap());
 
         if attach(pid).is_ok() {
-            pid
+            Ok(pid)
         } else {
-            panic!("Unable to attach to process {pid}");
+            Err(format!("Unable to attach to process `{pid}`"))
         }
     } else {
-        panic!("You must define a command or a PID to attach");
+        Err(format!("You must define a command or a PID to attach"))
     };
 
-    if !args.no_tui {
-        let mut ui = UI::new();
+    match process {
+        Ok(pid) => {
+            if !args.no_tui {
+                let mut ui = UI::new();
 
-        ui.start(pid, &args)?;
-    } else {
-        trace(pid, &args)?;
-    }
+                ui.start(pid, &args)?;
+            } else {
+                trace(pid, &args)?;
+            }
+        }
+        Err(e) => {
+            eprintln!("{}", e.red());
+        }
+    };
 
     Ok(())
 }
-- 
cgit v1.2.3-18-g5258