feat: make tmux template work with fish shell (#77)

0xricksanchez · web-flow · commit 18ec78b59569 · 2025-05-22T15:23:14.000+04:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -3,7 +3,7 @@ name = "afl_runner"
 authors = ["0x434b <mail@0x434b.dev"]
 repository = "https://github.com/0xricksanchez/AFL_Runner"
 description = "Scaling best-practice AFLPlusPlus fuzzing campaigns made easy"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2024"
 keywords = ["afl", "fuzzing", "fuzz-testing", "testing", "fuzzer"]
 license = "Apache-2.0"
@@ -22,22 +22,22 @@ path = "src/bin/generate_completions.rs"
 required-features = ["completion"]
 
 [dependencies]
-anyhow = "1.0.97"
-clap = { version = "4.5.32", features = ["wrap_help", "string", "derive"] }
-rand = "0.9.0"
-sysinfo = "0.33.1"
+anyhow = "1.0.98"
+clap = { version = "4.5.38", features = ["wrap_help", "string", "derive"] }
+rand = "0.9.1"
+sysinfo = "0.35.1"
 upon = "0.9.0"
-toml = "0.8.20"
+toml = "0.8.22"
 serde = { version = "1.0.219", features = ["derive"] }
 ratatui = { version = "0.29.0", features = ["crossterm", "all-widgets"] }
-crossterm = "0.28.1"
-chrono = "0.4.40"
+crossterm = "0.29.0"
+chrono = "0.4.41"
 uuid = { version = "1.16.0", features = ["v4"] }
-tempfile = "3.19.0"
-once_cell = "1.21.1"
+tempfile = "3.20.0"
+once_cell = "1.21.3"
 rayon = "1.10.0"
 glob = "0.3.2"
-clap_complete = { version = "4.5.46", optional = true }
+clap_complete = { version = "4.5.50", optional = true }
 
 [dev-dependencies]
 cargo-make = "0.37.24"
diff --git a/src/afl/cmd_gen.rs b/src/afl/cmd_gen.rs
@@ -55,7 +55,7 @@ impl AFLCmdGenerator {
             .collect::<Vec<String>>();
         if !gl_afl_env.is_empty() {
             println!(
-                "[!] Warning: Exported AFL++ environment variables found... Check generated commands!"
+                "[!] Warning: Exported AFL++ environment variables found {gl_afl_env:?}... Check generated commands!",
             );
         }
         gl_afl_env
diff --git a/src/afl/coverage.rs b/src/afl/coverage.rs
@@ -1,4 +1,4 @@
-use anyhow::{bail, Context, Result};
+use anyhow::{Context, Result, bail};
 use glob::glob;
 use rayon::prelude::*;
 use std::{
@@ -520,8 +520,7 @@ impl CoverageCollector {
             .par_chunks(1000)
             .enumerate()
             .map(|(i, chunk)| {
-                let temp_output = temp_dir.path().join(format!("temp_merged_{i}.profdata"));
-                
+                let temp_output = temp_dir.path().join(format!("temp_merged_{i}.profdata"));                
                 let output = Command::new("llvm-profdata")
                     .arg("merge")
                     .arg("-sparse")
diff --git a/src/afl/env.rs b/src/afl/env.rs
@@ -351,9 +351,10 @@ mod tests {
         assert_eq!(disable_trim_count, 2); // 60% of 4 rounded down = 2
 
         // Test ImportFirst (should be applied to all since runners < 8)
-        assert!(envs
-            .iter()
-            .all(|env| env.flags.contains(&AFLFlag::ImportFirst)));
+        assert!(
+            envs.iter()
+                .all(|env| env.flags.contains(&AFLFlag::ImportFirst))
+        );
     }
 
     #[test]
@@ -364,10 +365,11 @@ mod tests {
 
         assert!(envs.iter().take(0).all(|env| env.flags.len() == 1));
         // Check that the main fuzzer has at least the FINAL_SYNC flag set
-        assert!(envs
-            .iter()
-            .take(0)
-            .all(|env| env.flags.contains(&AFLFlag::FinalSync)));
+        assert!(
+            envs.iter()
+                .take(0)
+                .all(|env| env.flags.contains(&AFLFlag::FinalSync))
+        );
     }
 
     #[test]
@@ -376,9 +378,11 @@ mod tests {
         let envs = AFLEnv::new(Mode::MultipleCores, 20_u32, None, &mut rng);
 
         // Test that ImportFirst is not applied when runners >= 16
-        assert!(!envs
-            .iter()
-            .any(|env| env.flags.contains(&AFLFlag::ImportFirst)));
+        assert!(
+            !envs
+                .iter()
+                .any(|env| env.flags.contains(&AFLFlag::ImportFirst))
+        );
     }
 
     #[test]
diff --git a/src/argument_aggregator.rs b/src/argument_aggregator.rs
@@ -1,5 +1,5 @@
 use crate::cli::{AddSeedArgs, ArgMerge, Args, CovArgs, GenArgs, RunArgs};
-use anyhow::{bail, Context, Result};
+use anyhow::{Context, Result, bail};
 use std::{env, fs, path::PathBuf};
 
 static DEFAULT_AFL_CONFIG: &str = "aflr_cfg.toml";
diff --git a/src/commands/render_tui.rs b/src/commands/render_tui.rs
@@ -1,10 +1,10 @@
-use anyhow::{bail, Context, Result};
+use anyhow::{Context, Result, bail};
 use std::path::Path;
 
 use crate::{
     cli::TuiArgs,
     commands::Command,
-    tui::{session::CampaignData, Tui},
+    tui::{Tui, session::CampaignData},
 };
 
 pub struct RenderCommand<'a> {
diff --git a/src/runners/runner.rs b/src/runners/runner.rs
@@ -1,4 +1,4 @@
-use anyhow::{anyhow, Context, Result};
+use anyhow::{Context, Result, anyhow};
 use std::fs;
 use std::io::Write;
 use std::os::unix::fs::PermissionsExt;
@@ -8,7 +8,7 @@ use std::thread;
 use std::time::Duration;
 use tempfile::NamedTempFile;
 
-use crate::tui::{session::CampaignData, Tui};
+use crate::tui::{Tui, session::CampaignData};
 use crate::utils::system::{get_user_input, mkdir_helper};
 
 /// Template files for different session managers
diff --git a/src/runners/screen.rs b/src/runners/screen.rs
@@ -1,6 +1,6 @@
 use std::process::Command;
 
-use crate::runners::runner::{templates, Session, SessionManager};
+use crate::runners::runner::{Session, SessionManager, templates};
 
 /// Screen session manager implementation
 pub struct Screen;
diff --git a/src/runners/tmux.rs b/src/runners/tmux.rs
@@ -1,7 +1,7 @@
 use anyhow::{Context, Result};
 use std::process::Command;
 
-use crate::runners::runner::{templates, Session, SessionManager};
+use crate::runners::runner::{Session, SessionManager, templates};
 
 /// Tmux session manager implementation
 pub struct Tmux;
diff --git a/src/templates/tmux.txt b/src/templates/tmux.txt
@@ -1,9 +1,6 @@
 #!/bin/bash
 set -eu
-set +o history # Temporarily turn off history
-
-# Set HISTCONTROL to ignorespace
-HISTCONTROL="ignorespace"
+# HISTCONTROL="ignorespace" # For this script itself
 
 # Define the session name
 SESSION_NAME="{{ session_name }}"
@@ -18,95 +15,171 @@ COMMANDS=(
 )
 
 # Check if session does not exist yet
-tmux has-session -t $SESSION_NAME 2>/dev/null && {
+if tmux has-session -t "$SESSION_NAME" 2>/dev/null; then
     echo "Session '$SESSION_NAME' already exists. Aborting." >&2
     exit 1
-}
+fi
 
 # Start a new tmux session (detached)
-tmux -2u new-session -d -s $SESSION_NAME
+tmux -2u new-session -d -s "$SESSION_NAME"
 
 # Initialize the PID list
 AFLR_PID_LIST=""
 
-# Function to wait for shell to be ready
+# Function to wait for shell to be ready in a pane
 wait_for_shell() {
-    local window=$1
-    local max_attempts=10
+    local target_spec="$1" # e.g., session_name:window_name
+    local max_attempts=30  # 3 seconds total
     local attempt=1
-    local marker="AFLR_SHELL_READY_$$"  # Use PID to make marker unique
-    
-    # Send echo command with our marker
-    tmux send-keys -t "$SESSION_NAME:$window" " echo '$marker'" C-m
-    
-    while [ $attempt -le $max_attempts ]; do
-        if tmux capture-pane -p -t "$SESSION_NAME:$window" | grep -q "$marker"; then
-            sleep 0.1  # Small delay to ensure shell is fully ready after marker
+    local marker="AFLR_SHELL_READY_${$}_${RANDOM}"
+
+    # Send echo command with our marker. Leading space for history.
+    tmux send-keys -t "$target_spec" " echo '$marker'" C-m
+
+    while [ "$attempt" -le "$max_attempts" ]; do
+        if tmux capture-pane -p -t "$target_spec" | grep -qF "$marker"; then
+            sleep 0.1 # Small delay to ensure shell is fully ready after marker
             return 0
         fi
         sleep 0.1
         attempt=$((attempt + 1))
     done
+    echo "Error: Shell in pane '$target_spec' did not become ready (marker: $marker not found)." >&2
+    # For debugging, uncomment below:
+    # echo "DEBUG: Pane content for $target_spec on wait_for_shell timeout:" >&2
+    # tmux capture-pane -p -t "$target_spec" >&2
     return 1
 }
 
+# Function to get shell type in a pane using tmux's knowledge
+get_pane_shell_type() {
+    local target_spec="$1"
+    local pane_cmd_output
+    local shell_name_tmp
+    # Max attempts for display-message to return non-empty. Should be quick.
+    local max_attempts=10 # 10 * 0.1s = 1 second.
+    local attempt=1
+
+    while [ "$attempt" -le "$max_attempts" ]; do
+        # Get the command running in the pane directly from tmux
+        pane_cmd_output=$(tmux display-message -p -t "$target_spec" '#{pane_current_command}')
+        
+        if [ -n "$pane_cmd_output" ]; then
+            # Extract the basename of the command (e.g., /usr/bin/bash -> bash)
+            shell_name_tmp=$(basename "$pane_cmd_output")
+
+            # Uncomment for debugging what tmux reports as the pane command:
+            # echo "DEBUG: For target '$target_spec', pane_current_command is '$pane_cmd_output', basename is '$shell_name_tmp'." >&2
+
+            if [[ "$shell_name_tmp" == *bash* ]]; then echo "bash"; return 0; fi
+            if [[ "$shell_name_tmp" == *fish* ]]; then echo "fish"; return 0; fi
+            if [[ "$shell_name_tmp" == *zsh* ]]; then echo "zsh"; return 0; fi
+            # Covers 'sh', 'dash', and other simple or versioned names like 'bash-5.1'
+            if [[ "$shell_name_tmp" =~ ^[a-zA-Z0-9._-]+$ ]]; then
+                echo "$shell_name_tmp" # Return the detected name
+                return 0
+            fi
+            
+            # If output was received but not classified by specific checks above
+            echo "Warning: Unclassified pane command '$pane_cmd_output' (basename '$shell_name_tmp') for target '$target_spec'. Defaulting to 'bash'." >&2
+            echo "bash" # Default to bash
+            return 1 # Indicate it's a fallback/guess
+        fi
+        sleep 0.1 # Wait a bit if pane_current_command is not yet set (e.g., pane initializing)
+        attempt=$((attempt + 1))
+    done
+    
+    echo "Warning: Could not determine shell for '$target_spec' (pane_current_command was empty after $max_attempts attempts). Defaulting to 'bash'." >&2
+    echo "bash" 
+    return 1 # Indicate failure / fallback
+}
+
+
 # Create and rename a window for each command
 for i in "${!COMMANDS[@]}"; do
     WINDOW_NAME="window-$i"
-    TEMP_PID_FILE="/tmp/aflr_pid_${i}.txt"
-    
-    if [ $i -eq 0 ]; then
-        # For the first command, send it to the first window and rename it
-        tmux rename-window -t $SESSION_NAME $WINDOW_NAME
-        wait_for_shell "$WINDOW_NAME"
-        tmux send-keys -t $SESSION_NAME " { ${COMMANDS[$i]} & echo \$! > $TEMP_PID_FILE; clear; fg; }" C-m
+    TEMP_PID_FILE="/tmp/aflr_pid_${SESSION_NAME}_${i}_$$.txt"
+    rm -f "$TEMP_PID_FILE" # Clean up any old temp file
+
+    current_target_spec=""
+
+    if [ "$i" -eq 0 ]; then
+        tmux rename-window -t "$SESSION_NAME" "$WINDOW_NAME"
+        current_target_spec="$SESSION_NAME:$WINDOW_NAME"
     else
-        # For subsequent commands, create new windows and rename them
-        tmux new-window -t $SESSION_NAME -n $WINDOW_NAME
-        wait_for_shell "$WINDOW_NAME"
-        tmux send-keys -t $SESSION_NAME:$WINDOW_NAME " { ${COMMANDS[$i]} & echo \$! > $TEMP_PID_FILE; clear; fg; }" C-m
+        tmux new-window -t "$SESSION_NAME" -n "$WINDOW_NAME"
+        current_target_spec="$SESSION_NAME:$WINDOW_NAME"
+    fi
+
+    if ! wait_for_shell "$current_target_spec"; then
+        echo "Aborting due to shell readiness failure in $current_target_spec." >&2
+        tmux kill-session -t "$SESSION_NAME" 2>/dev/null || true
+        exit 1
+    fi
+
+    # Clear the marker left by wait_for_shell from the pane display
+    tmux send-keys -t "$current_target_spec" " clear" C-m
+    sleep 0.2 # Give clear a moment to execute
+
+    # Detect the shell type.
+    _detected_shell=$(get_pane_shell_type "$current_target_spec")
+    
+    # Uncomment for debugging the detected shell:
+    # echo "INFO: Detected shell in $current_target_spec as '$_detected_shell'."
+
+    _command_to_send_in_pane=""
+    # Ensure commands sent to panes start with a space to try and avoid shell history
+    if [[ "$_detected_shell" == "fish" ]]; then
+        _command_to_send_in_pane=" ${COMMANDS[$i]} &; set TMUX_SCRIPT_LAST_PID (jobs -q --last --pid); echo \$TMUX_SCRIPT_LAST_PID > \"$TEMP_PID_FILE\"; clear; fg %1 2>/dev/null || true"
+    else # Bash, Zsh, sh, dash, etc. (POSIX-like $!)
+        _command_to_send_in_pane=" { ${COMMANDS[$i]} & echo \$! > \"$TEMP_PID_FILE\"; clear; fg %1 2>/dev/null || true; }"
     fi
     
-    # Add a small delay to ensure the command has started
+    tmux send-keys -t "$current_target_spec" "$_command_to_send_in_pane" C-m
+    
+    # Add a small delay for the command to start/background before piping output
+    sleep 0.1 
+    tmux pipe-pane -o -t "$current_target_spec" "cat >> \"$LOG_FILE\""
+    
+    # Delay before processing the next window
     sleep 0.2
 done
 
-# Wait for all commands to start and PIDs to be written
-sleep 1
+echo "Waiting for PIDs to be written..."
+sleep 1.5 # Increased wait time for all commands to start and write PIDs
 
 # Capture the PIDs and verify processes are running
 for i in "${!COMMANDS[@]}"; do
-    TEMP_PID_FILE="/tmp/aflr_pid_${i}.txt"
+    TEMP_PID_FILE="/tmp/aflr_pid_${SESSION_NAME}_${i}_$$.txt"
+    PID_VALUE=""
+
     if [ -f "$TEMP_PID_FILE" ]; then
-        PID=$(cat "$TEMP_PID_FILE")
-        if [ -n "$PID" ]; then
-            # Check if the process is actually running
-            if kill -0 "$PID" 2>/dev/null; then
+        PID_VALUE=$(cat "$TEMP_PID_FILE")
+        if [[ "$PID_VALUE" =~ ^[0-9]+$ ]]; then # Validate PID is a number
+            if kill -0 "$PID_VALUE" 2>/dev/null; then # Check if process is running
                 if [ -z "$AFLR_PID_LIST" ]; then
-                    AFLR_PID_LIST="$PID"
+                    AFLR_PID_LIST="$PID_VALUE"
                 else
-                    AFLR_PID_LIST="$AFLR_PID_LIST:$PID"
+                    AFLR_PID_LIST="$AFLR_PID_LIST:$PID_VALUE"
                 fi
             else
-                echo "Process $PID (command: ${COMMANDS[$i]}) failed to start or crashed immediately!" >&2
-                exit 1
+                echo "Process $PID_VALUE (command: ${COMMANDS[$i]}) not running or crashed. Check $LOG_FILE." >&2
             fi
+        else
+            echo "Invalid PID '$PID_VALUE' found in $TEMP_PID_FILE for command: ${COMMANDS[$i]}. Likely a shell detection or command execution issue. Check $LOG_FILE." >&2
         fi
         rm "$TEMP_PID_FILE"
     else
-        echo "PID file not found for command: ${COMMANDS[$i]}" >&2
-        exit 1
+        echo "PID file $TEMP_PID_FILE not found for command: ${COMMANDS[$i]}. Process likely failed to start correctly. Check $LOG_FILE." >&2
     fi
 done
 
-# Redirect tmux server log to a specific file
-tmux pipe-pane -o -t $SESSION_NAME "cat >> $LOG_FILE"
-
 if [ -z "$AFLR_PID_LIST" ]; then
-    echo "Failed to start AFLR session '$SESSION_NAME'" >&2
+    echo "Failed to start any AFLR processes in session '$SESSION_NAME'. Check $LOG_FILE." >&2
+    tmux kill-session -t "$SESSION_NAME" 2>/dev/null || true # Clean up session
     exit 1
 else
-    echo $AFLR_PID_LIST > "$PID_FILE"
+    echo "$AFLR_PID_LIST" > "$PID_FILE"
+    echo "AFLR session '$SESSION_NAME' started with PIDs: $AFLR_PID_LIST. Log: $LOG_FILE. PID file: $PID_FILE"
+    echo "Attach with: tmux attach -t $SESSION_NAME"
 fi
-
-set -o history # Turn history back on

Original file line number	Diff line number	Diff line change
`@@ -55,7 +55,7 @@ impl AFLCmdGenerator {`
`55`	`55`	`.collect::<Vec<String>>();`
`56`	`56`	`if !gl_afl_env.is_empty() {`
`57`	`57`	`println!(`
`58`		`- "[!] Warning: Exported AFL++ environment variables found... Check generated commands!"`
	`58`	`+ "[!] Warning: Exported AFL++ environment variables found {gl_afl_env:?}... Check generated commands!",`
`59`	`59`	`);`
`60`	`60`	`}`
`61`	`61`	`gl_afl_env`