Add SSH agent key selection and fix MapReduce outfile handling

This commit adds two major features and fixes:

1. SSH Agent Key Selection:
   - Add --agentKeyIndex flag to select specific SSH agent key (0-based)
   - Solves "too many authentication failures" with multiple SSH keys
   - Default -1 uses all keys (backwards compatible)
   - Available in dtail, dcat, dgrep, dmap commands

2. MapReduce Outfile Fixes:
   - CSV files now written at every interval, not just on exit
   - Proper signal handling (SIGTERM/SIGINT) with graceful shutdown
   - 5-second grace period for cleanup before force exit
   - Fixes issue where outfile remained as .tmp during execution

Usage:
  dtail --servers host --agentKeyIndex 0 --query '...' outfile results.csv

This is particularly useful with YubiKey/hardware tokens where many
keys are loaded in the SSH agent, and for monitoring MapReduce results
in real-time as they're computed.

Co-authored-by: Cursor <cursoragent@cursor.com>

1 files changed, 45 insertions, 0 deletions

diff --git a/internal/io/signal/signal.go b/internal/io/signal/signal.go
index c01e82a..e94de42 100644
--- a/internal/io/signal/signal.go
+++ b/internal/io/signal/signal.go
@@ -10,7 +10,52 @@ import (
 	"github.com/mimecast/dtail/internal/config"
 )
 
+// InterruptChWithCancel returns a channel for "please print stats" signalling.
+// It accepts a cancel function to properly shutdown when termination signals are received.
+func InterruptChWithCancel(ctx context.Context, cancel context.CancelFunc) <-chan string {
+	sigIntCh := make(chan os.Signal, 10)
+	gosignal.Notify(sigIntCh, os.Interrupt)
+	sigOtherCh := make(chan os.Signal, 10)
+	gosignal.Notify(sigOtherCh, syscall.SIGHUP, syscall.SIGTERM, syscall.SIGQUIT)
+	statsCh := make(chan string)
+
+	go func() {
+		for {
+			select {
+			case <-sigIntCh:
+				select {
+				case statsCh <- "Hint: Hit Ctrl+C again to exit":
+					select {
+					case <-sigIntCh:
+						cancel()
+						// Wait longer to allow MapReduce cleanup, then force exit if still running
+						go func() {
+							time.Sleep(5 * time.Second)
+							os.Exit(0)
+						}()
+					case <-time.After(time.Second * time.Duration(config.InterruptTimeoutS)):
+					}
+				default:
+					// Stats already printed.
+				}
+			case <-sigOtherCh:
+				// Cancel context to allow graceful shutdown (MapReduce outfile cleanup, etc.)
+				cancel()
+				// Wait longer to allow MapReduce cleanup, then force exit if still running
+				go func() {
+					time.Sleep(5 * time.Second)
+					os.Exit(0)
+				}()
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+	return statsCh
+}
+
 // InterruptCh returns a channel for "please print stats" signalling.
+// Deprecated: Use InterruptChWithCancel for proper cleanup on termination signals.
 func InterruptCh(ctx context.Context) <-chan string {
 	sigIntCh := make(chan os.Signal, 10)
 	gosignal.Notify(sigIntCh, os.Interrupt)