From 10691b461ff73b11ce23d217c63d4f70bc916982 Mon Sep 17 00:00:00 2001
From: c_sto <7466346+C-Sto@users.noreply.github.com>
Date: Sun, 4 Nov 2018 18:53:07 +0800
Subject: [PATCH] add ability to dynamically adjust workers (!!!)

---
 librecursebuster/logic.go   | 12 +++++++++---
 librecursebuster/output.go  | 16 +++++++++++-----
 librecursebuster/structs.go | 21 ++++++++++++---------
 librecursebuster/ui.go      | 31 ++++++++++++++++++++++++++++++-
 main.go                     |  2 +-
 5 files changed, 63 insertions(+), 19 deletions(-)

diff --git a/librecursebuster/logic.go b/librecursebuster/logic.go
index 06191e8..6ecd77d 100644
--- a/librecursebuster/logic.go
+++ b/librecursebuster/logic.go
@@ -135,8 +135,12 @@ type workUnit struct {
 
 func (gState *State) testWorker() {
 	for {
-		w := <-gState.Chans.workersChan
-		gState.testURL(w.Method, w.URLString, gState.Client)
+		select {
+		case w := <-gState.Chans.workersChan:
+			gState.testURL(w.Method, w.URLString, gState.Client)
+		case <-gState.Chans.lessWorkersChan:
+			return
+		}
 	}
 }
 
@@ -144,6 +148,9 @@ func (gState *State) testURL(method string, urlString string, client *http.Clien
 	defer func() {
 		gState.wg.Done()
 		atomic.AddUint64(gState.TotalTested, 1)
+		if gState.DirbProgress != nil { //shoudl probably check if there is a wordlist, but this will do..
+			atomic.AddUint32(gState.DirbProgress, 1)
+		}
 	}()
 	select {
 	case gState.Chans.testChan <- method + ":" + urlString:
@@ -285,7 +292,6 @@ func (gState *State) dirBust(page SpiderPage) {
 				gState.Checked[method+page.URL+word] = true
 				gState.CMut.Unlock()
 				//if gState.Cfg.MaxDirs == 1 {
-				atomic.AddUint32(gState.DirbProgress, 1)
 				//}
 			}
 		}
diff --git a/librecursebuster/output.go b/librecursebuster/output.go
index 2de63a2..b089cc1 100644
--- a/librecursebuster/output.go
+++ b/librecursebuster/output.go
@@ -166,16 +166,21 @@ func (gState *State) writeStatus(s string) {
 		// handle error
 	}
 	v.Clear()
-	fmt.Fprintln(v, gState.getStatus())
+	//fmt.Fprintln(v, gState.getStatus())
+	timeFormat := "%s (+%vs) " + gState.getStatus()
+	fmt.Fprintln(v, fmt.Sprintf(timeFormat, time.Now().Format("2006-01-02 15:04:05"), int(time.Since(gState.StartTime).Seconds())))
 	sprint := ""
 	if len(gState.WordList) > 0 {
-		sprint = fmt.Sprintf("[%.2f%%%%]%s", 100*float64(atomic.LoadUint32(gState.DirbProgress))/float64(len(gState.WordList)), s)
+		sprint = fmt.Sprintf("[%.2f%%]%s", 100*float64(atomic.LoadUint32(gState.DirbProgress))/float64(len(gState.WordList)), s)
 	} else {
 		sprint = fmt.Sprintf("Waiting on %v items", gState.wg)
 	}
+	fmt.Fprintln(v, "ctrl + (c) Quit, (x) Stop current dir, (arrow up|down) Move one line, (pgup|pgdown) Move 10 lines")
+	fmt.Fprintln(v, "ctrl + (t) Add worker, (y) Remove worker")
 	fmt.Fprintln(v, sprint)
-	fmt.Fprintln(v, "ctrl + [(c) quit, (x) stop current dir], (arrow up/down) move one line, (pgup/pgdown) move 10 lines")
-	fmt.Fprintln(v, time.Now().String())
+	//Time format: yyyy-mm-dd hh:mm:ss tz (elapsed seconds)
+	//2018-11-04 18:13:40.6721974 +0800 AWST m=+4.232677701
+	//time.Now().cl
 	return //nil
 }
 
@@ -245,7 +250,8 @@ func (gState *State) StatusPrinter() {
 }
 
 func (gState *State) getStatus() string {
-	return fmt.Sprintf("Tested: %d Speed(2s): %d/s Speed: %d/s",
+	return fmt.Sprintf("Workers: %d Tested: %d Speed(2s): %d/s Speed: %d/s",
+		atomic.LoadUint32(gState.workerCount),
 		atomic.LoadUint64(gState.TotalTested),
 		atomic.LoadUint64(gState.PerSecondShort),
 		atomic.LoadUint64(gState.PerSecondLong),
diff --git a/librecursebuster/structs.go b/librecursebuster/structs.go
index ef730a8..072a86e 100644
--- a/librecursebuster/structs.go
+++ b/librecursebuster/structs.go
@@ -113,10 +113,10 @@ type chans struct {
 	newPagesChan,
 	confirmedChan chan SpiderPage
 
-	//workersChan chan struct{}
-	workersChan chan workUnit
-	printChan   chan OutLine
-	testChan    chan string
+	workersChan     chan workUnit
+	lessWorkersChan chan struct{}
+	printChan       chan OutLine
+	testChan        chan string
 }
 
 func (c *chans) GetWorkers() chan workUnit {
@@ -125,10 +125,11 @@ func (c *chans) GetWorkers() chan workUnit {
 
 func (chans) Init() *chans {
 	return &chans{
-		pagesChan:     make(chan SpiderPage, 1000),
-		newPagesChan:  make(chan SpiderPage, 10000),
-		confirmedChan: make(chan SpiderPage, 1000),
-		workersChan:   make(chan workUnit, 1000),
+		pagesChan:       make(chan SpiderPage, 1000),
+		newPagesChan:    make(chan SpiderPage, 10000),
+		confirmedChan:   make(chan SpiderPage, 1000),
+		workersChan:     make(chan workUnit, 1000),
+		lessWorkersChan: make(chan struct{}, 5), //too bad if you want to add more than 5 at a time ok
 		//maxDirs := make(chan struct{}, cfg.MaxDirs),
 		testChan:  make(chan string, 100),
 		printChan: make(chan OutLine, 100),
@@ -146,6 +147,7 @@ type State struct {
 	TotalTested    *uint64
 	PerSecondShort *uint64 //how many tested over 2 seconds or so
 	PerSecondLong  *uint64
+	workerCount    *uint32 //probably doesn't need to be async safe, but whatever
 	StartTime      time.Time
 	Blacklist      map[string]bool
 	Whitelist      map[string]bool
@@ -398,7 +400,8 @@ func (gState *State) SetupState() {
 			//atomic.AddUint32(gState.WordlistLen, 1)
 		}
 	}
-
+	workers := uint32(gState.Cfg.Threads)
+	gState.workerCount = &workers
 	gState.StartTime = time.Now()
 	gState.PerSecondShort = new(uint64)
 	gState.PerSecondLong = new(uint64)
diff --git a/librecursebuster/ui.go b/librecursebuster/ui.go
index 6b8c94b..9371d89 100644
--- a/librecursebuster/ui.go
+++ b/librecursebuster/ui.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	ui "github.com/jroimartin/gocui"
 )
@@ -46,7 +47,18 @@ func (s *State) StartUI(uiWG *sync.WaitGroup, quitChan chan struct{}) {
 	err = s.ui.SetKeybinding("", ui.KeyArrowDown, ui.ModNone, scrollDown)
 	if err != nil {
 		panic(err)
-	} /* Mouse stuff broke copying out of the terminal... not ideal
+	}
+
+	err = s.ui.SetKeybinding("", ui.KeyCtrlT, ui.ModNone, s.addWorker)
+	if err != nil {
+		panic(err)
+	}
+
+	err = s.ui.SetKeybinding("", ui.KeyCtrlY, ui.ModNone, s.stopWorker) //wtf? no shift modifier??
+	if err != nil {
+		panic(err)
+	}
+	/* Mouse stuff broke copying out of the terminal... not ideal
 	err = s.ui.SetKeybinding("", ui.MouseWheelUp, ui.ModNone, scrollUp)
 	if err != nil {
 		panic(err)
@@ -62,6 +74,23 @@ func (s *State) StartUI(uiWG *sync.WaitGroup, quitChan chan struct{}) {
 	}
 }
 
+func (gState *State) addWorker(g *ui.Gui, v *ui.View) error {
+	atomic.AddUint32(gState.workerCount, 1)
+	go gState.testWorker()
+	return nil
+}
+
+func (gState *State) stopWorker(g *ui.Gui, v *ui.View) error {
+	count := atomic.LoadUint32(gState.workerCount)
+	if count == 0 { //avoid underflow
+		return nil
+	}
+	count = count - 1
+	atomic.StoreUint32(gState.workerCount, count)
+	gState.Chans.lessWorkersChan <- struct{}{}
+	return nil
+}
+
 //StopUI should be called when closing the program. It prints out the lines in the main view buffer to stdout, and closes the ui object
 func (gState *State) StopUI() {
 	p, _ := gState.ui.View("Main")
diff --git a/main.go b/main.go
index bf3b923..061b11c 100644
--- a/main.go
+++ b/main.go
@@ -14,7 +14,7 @@ import (
 	"github.com/fatih/color"
 )
 
-const version = "1.5.17"
+const version = "1.6.0"
 
 func main() {
 	if runtime.GOOS == "windows" { //lol goos