Skip to content

Commit

Permalink
add ability to dynamically adjust workers (!!!)
Browse files Browse the repository at this point in the history
  • Loading branch information
C-Sto committed Nov 4, 2018
1 parent 784084d commit 10691b4
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 19 deletions.
12 changes: 9 additions & 3 deletions librecursebuster/logic.go
Original file line number Diff line number Diff line change
Expand Up @@ -135,15 +135,22 @@ 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
}
}
}

func (gState *State) testURL(method string, urlString string, client *http.Client) {
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:
Expand Down Expand Up @@ -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)
//}
}
}
Expand Down
16 changes: 11 additions & 5 deletions librecursebuster/output.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down Expand Up @@ -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),
Expand Down
21 changes: 12 additions & 9 deletions librecursebuster/structs.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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),
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down
31 changes: 30 additions & 1 deletion librecursebuster/ui.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"fmt"
"strings"
"sync"
"sync/atomic"

ui "github.com/jroimartin/gocui"
)
Expand Down Expand Up @@ -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)
Expand All @@ -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")
Expand Down
2 changes: 1 addition & 1 deletion main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 10691b4

Please sign in to comment.