From e89b2db204cfcc6fafc33f83206dcaf97025180e Mon Sep 17 00:00:00 2001 From: Sheikh Shaheer Imran Date: Thu, 1 Jan 2026 02:36:09 +0100 Subject: [PATCH] Add pause/resume support with SPACE key - Add PauseController for worker synchronization - Add keyboard listener for SPACE key toggle - Show [PAUSED] in progress bar when paused --- cli/const_windows.go | 2 +- cli/gobuster.go | 17 +++- cli/keyboard.go | 85 +++++++++++++++++++ cli/keyboard_windows.go | 83 ++++++++++++++++++ libgobuster/libgobuster.go | 7 ++ libgobuster/pause.go | 85 +++++++++++++++++++ libgobuster/pause_test.go | 170 +++++++++++++++++++++++++++++++++++++ 7 files changed, 447 insertions(+), 2 deletions(-) create mode 100644 cli/keyboard.go create mode 100644 cli/keyboard_windows.go create mode 100644 libgobuster/pause.go create mode 100644 libgobuster/pause_test.go diff --git a/cli/const_windows.go b/cli/const_windows.go index 1224bb69..b1fe29ba 100644 --- a/cli/const_windows.go +++ b/cli/const_windows.go @@ -3,5 +3,5 @@ package cli const ( - TerminalClearLine = "\r\r" + TerminalClearLine = "\r\x1b[2K" ) diff --git a/cli/gobuster.go b/cli/gobuster.go index 3841db39..e97807fe 100644 --- a/cli/gobuster.go +++ b/cli/gobuster.go @@ -102,7 +102,12 @@ func printProgress(g *libgobuster.Gobuster) { if math.IsNaN(float64(percent)) { percent = 0.0 } - s := fmt.Sprintf("%sProgress: %d / %d (%3.2f%%)", TerminalClearLine, requestsIssued, requestsExpected, percent) + var s string + if g.Pause.IsPaused() { + s = fmt.Sprintf("%s[PAUSED] Progress: %d / %d (%3.2f%%) - Press SPACE to resume", TerminalClearLine, requestsIssued, requestsExpected, percent) + } else { + s = fmt.Sprintf("%sProgress: %d / %d (%3.2f%%)", TerminalClearLine, requestsIssued, requestsExpected, percent) + } _, _ = fmt.Fprint(os.Stderr, s) } @@ -180,6 +185,16 @@ func Gobuster(ctx context.Context, opts *libgobuster.Options, plugin libgobuster opts.NoProgress = true } + // Start keyboard listener for pause/resume (only in TTY mode and not reading from stdin) + if !opts.NoProgress && opts.Wordlist != "-" { + if !opts.Quiet { + log.Println("[*] Press SPACE to pause") + log.Println(ruler) + } + cleanup := StartKeyboardListener(ctxCancel, gobuster, cancel) + defer cleanup() + } + // our waitgroup for all goroutines // this ensures all goroutines are finished // when we call wg.Wait() diff --git a/cli/keyboard.go b/cli/keyboard.go new file mode 100644 index 00000000..9268e169 --- /dev/null +++ b/cli/keyboard.go @@ -0,0 +1,85 @@ +//go:build !windows + +package cli + +import ( + "context" + "os" + "os/signal" + "sync" + "syscall" + + "github.com/OJ/gobuster/v3/libgobuster" + "golang.org/x/term" +) + +func StartKeyboardListener(ctx context.Context, g *libgobuster.Gobuster, cancel context.CancelFunc) func() { + fd := int(os.Stdin.Fd()) + oldState, err := term.MakeRaw(fd) + if err != nil { + return func() {} + } + + var restoreOnce sync.Once + restoreTerminal := func() { + restoreOnce.Do(func() { + _ = term.Restore(fd, oldState) + }) + } + + // restore terminal on signals + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) + + go func() { + select { + case sig := <-sigCh: + restoreTerminal() + signal.Stop(sigCh) + // re-raise signal + p, _ := os.FindProcess(os.Getpid()) + _ = p.Signal(sig.(syscall.Signal)) + case <-ctx.Done(): + signal.Stop(sigCh) + } + }() + + // read keys and send to channel + keyCh := make(chan byte, 1) + go func() { + buf := make([]byte, 1) + for { + n, err := os.Stdin.Read(buf) + if err != nil || n == 0 { + return + } + select { + case keyCh <- buf[0]: + case <-ctx.Done(): + return + } + } + }() + + // handle key events + go func() { + for { + select { + case <-ctx.Done(): + return + case key := <-keyCh: + if key == ' ' { + g.Pause.Toggle() + } + // Ctrl+C + if key == 3 { + restoreTerminal() + cancel() + return + } + } + } + }() + + return restoreTerminal +} diff --git a/cli/keyboard_windows.go b/cli/keyboard_windows.go new file mode 100644 index 00000000..d459a86a --- /dev/null +++ b/cli/keyboard_windows.go @@ -0,0 +1,83 @@ +//go:build windows + +package cli + +import ( + "context" + "os" + "os/signal" + "sync" + "syscall" + + "github.com/OJ/gobuster/v3/libgobuster" + "golang.org/x/term" +) + +func StartKeyboardListener(ctx context.Context, g *libgobuster.Gobuster, cancel context.CancelFunc) func() { + fd := int(syscall.Stdin) + oldState, err := term.MakeRaw(fd) + if err != nil { + return func() {} + } + + var restoreOnce sync.Once + restoreTerminal := func() { + restoreOnce.Do(func() { + _ = term.Restore(fd, oldState) + }) + } + + // restore terminal on signals + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) + + go func() { + select { + case <-sigCh: + restoreTerminal() + signal.Stop(sigCh) + cancel() + case <-ctx.Done(): + signal.Stop(sigCh) + } + }() + + // read keys and send to channel + keyCh := make(chan byte, 1) + go func() { + buf := make([]byte, 1) + for { + n, err := os.Stdin.Read(buf) + if err != nil || n == 0 { + return + } + select { + case keyCh <- buf[0]: + case <-ctx.Done(): + return + } + } + }() + + // handle key events + go func() { + for { + select { + case <-ctx.Done(): + return + case key := <-keyCh: + if key == ' ' { + g.Pause.Toggle() + } + // Ctrl+C + if key == 3 { + restoreTerminal() + cancel() + return + } + } + } + }() + + return restoreTerminal +} diff --git a/libgobuster/libgobuster.go b/libgobuster/libgobuster.go index 1e803d62..a0fcaee3 100644 --- a/libgobuster/libgobuster.go +++ b/libgobuster/libgobuster.go @@ -30,6 +30,7 @@ type Gobuster struct { Logger *Logger plugin GobusterPlugin Progress *Progress + Pause *PauseController } type Guess struct { @@ -50,6 +51,7 @@ func NewGobuster(opts *Options, plugin GobusterPlugin, logger *Logger) (*Gobuste g.plugin = plugin g.Logger = logger g.Progress = NewProgress() + g.Pause = NewPauseController() return &g, nil } @@ -57,6 +59,11 @@ func NewGobuster(opts *Options, plugin GobusterPlugin, logger *Logger) (*Gobuste func (g *Gobuster) worker(ctx context.Context, guessChan <-chan *Guess, successChan chan<- *Guess, wg *sync.WaitGroup) { defer wg.Done() for { + // Check if paused and wait for resume + if err := g.Pause.Wait(ctx); err != nil { + return + } + // Prioritize stopping when the context is done select { case <-ctx.Done(): diff --git a/libgobuster/pause.go b/libgobuster/pause.go new file mode 100644 index 00000000..cb0009da --- /dev/null +++ b/libgobuster/pause.go @@ -0,0 +1,85 @@ +package libgobuster + +import ( + "context" + "sync" +) + +// PauseController handles pause/resume for workers +type PauseController struct { + mu sync.RWMutex + paused bool + pauseCh chan struct{} + resumeCh chan struct{} +} + +// NewPauseController creates a new unpaused controller +func NewPauseController() *PauseController { + return &PauseController{ + pauseCh: make(chan struct{}), + resumeCh: make(chan struct{}), + } +} + +// Pause pauses the controller +func (p *PauseController) Pause() { + p.mu.Lock() + defer p.mu.Unlock() + if !p.paused { + p.paused = true + close(p.pauseCh) + } +} + +// Resume unpauses the controller +func (p *PauseController) Resume() { + p.mu.Lock() + defer p.mu.Unlock() + if p.paused { + p.paused = false + p.pauseCh = make(chan struct{}) + close(p.resumeCh) + p.resumeCh = make(chan struct{}) + } +} + +// Toggle flips pause state and returns the new state +func (p *PauseController) Toggle() bool { + p.mu.Lock() + defer p.mu.Unlock() + if p.paused { + p.paused = false + p.pauseCh = make(chan struct{}) + close(p.resumeCh) + p.resumeCh = make(chan struct{}) + } else { + p.paused = true + close(p.pauseCh) + } + return p.paused +} + +// IsPaused returns true if paused +func (p *PauseController) IsPaused() bool { + p.mu.RLock() + defer p.mu.RUnlock() + return p.paused +} + +// Wait blocks until resumed or context is cancelled +func (p *PauseController) Wait(ctx context.Context) error { + p.mu.RLock() + if !p.paused { + p.mu.RUnlock() + return nil + } + resumeCh := p.resumeCh + p.mu.RUnlock() + + select { + case <-ctx.Done(): + return ctx.Err() + case <-resumeCh: + return nil + } +} diff --git a/libgobuster/pause_test.go b/libgobuster/pause_test.go new file mode 100644 index 00000000..aef2a9c8 --- /dev/null +++ b/libgobuster/pause_test.go @@ -0,0 +1,170 @@ +package libgobuster + +import ( + "context" + "sync" + "testing" + "time" +) + +func TestNewPauseController(t *testing.T) { + t.Parallel() + pc := NewPauseController() + if pc == nil { + t.Fatal("NewPauseController returned nil") + } + if pc.IsPaused() { + t.Fatal("NewPauseController should start in unpaused state") + } +} + +func TestPauseControllerPause(t *testing.T) { + t.Parallel() + pc := NewPauseController() + + pc.Pause() + if !pc.IsPaused() { + t.Fatal("IsPaused should return true after Pause()") + } + + pc.Pause() + if !pc.IsPaused() { + t.Fatal("IsPaused should still return true after double Pause()") + } +} + +func TestPauseControllerResume(t *testing.T) { + t.Parallel() + pc := NewPauseController() + + pc.Pause() + pc.Resume() + if pc.IsPaused() { + t.Fatal("IsPaused should return false after Resume()") + } + + pc.Resume() + if pc.IsPaused() { + t.Fatal("IsPaused should still return false after double Resume()") + } +} + +func TestPauseControllerToggle(t *testing.T) { + t.Parallel() + pc := NewPauseController() + + paused := pc.Toggle() + if !paused { + t.Fatal("Toggle should return true when transitioning to paused") + } + if !pc.IsPaused() { + t.Fatal("IsPaused should return true after Toggle()") + } + + paused = pc.Toggle() + if paused { + t.Fatal("Toggle should return false when transitioning to unpaused") + } + if pc.IsPaused() { + t.Fatal("IsPaused should return false after second Toggle()") + } +} + +func TestPauseControllerWaitNotPaused(t *testing.T) { + t.Parallel() + pc := NewPauseController() + ctx := context.Background() + + start := time.Now() + err := pc.Wait(ctx) + elapsed := time.Since(start) + + if err != nil { + t.Fatalf("Wait should return nil when not paused, got: %v", err) + } + if elapsed > 10*time.Millisecond { + t.Fatalf("Wait should return immediately when not paused, took: %v", elapsed) + } +} + +func TestPauseControllerWaitPausedThenResume(t *testing.T) { + t.Parallel() + pc := NewPauseController() + ctx := context.Background() + + pc.Pause() + + var wg sync.WaitGroup + var waitErr error + var waitDuration time.Duration + + wg.Add(1) + go func() { + defer wg.Done() + start := time.Now() + waitErr = pc.Wait(ctx) + waitDuration = time.Since(start) + }() + + time.Sleep(50 * time.Millisecond) + pc.Resume() + wg.Wait() + + if waitErr != nil { + t.Fatalf("Wait should return nil after resume, got: %v", waitErr) + } + if waitDuration < 40*time.Millisecond { + t.Fatalf("Wait should have blocked for at least 40ms, took: %v", waitDuration) + } +} + +func TestPauseControllerWaitContextCancel(t *testing.T) { + t.Parallel() + pc := NewPauseController() + ctx, cancel := context.WithCancel(context.Background()) + + pc.Pause() + + var wg sync.WaitGroup + var waitErr error + var waitDuration time.Duration + + wg.Add(1) + go func() { + defer wg.Done() + start := time.Now() + waitErr = pc.Wait(ctx) + waitDuration = time.Since(start) + }() + + time.Sleep(50 * time.Millisecond) + cancel() + wg.Wait() + + if waitErr != context.Canceled { + t.Fatalf("Wait should return context.Canceled, got: %v", waitErr) + } + if waitDuration < 40*time.Millisecond { + t.Fatalf("Wait should have blocked for at least 40ms, took: %v", waitDuration) + } +} + +func TestPauseControllerConcurrency(t *testing.T) { + t.Parallel() + pc := NewPauseController() + + var wg sync.WaitGroup + for i := 0; i < 100; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 100; j++ { + pc.Pause() + _ = pc.IsPaused() + pc.Resume() + } + }() + } + + wg.Wait() +}