|
package main |
|
|
|
import ( |
|
"bufio" |
|
"context" |
|
"flag" |
|
"fmt" |
|
"log" |
|
"os" |
|
"os/exec" |
|
"runtime" |
|
"sync" |
|
"sync/atomic" |
|
) |
|
|
|
func runCmdReturnExitCode(ctx context.Context, name string, args ...string) (int, error) { |
|
cmd := exec.CommandContext(ctx, name, args...) |
|
if err := cmd.Run(); err == nil { |
|
if ps := cmd.ProcessState; ps != nil { |
|
return ps.ExitCode(), nil |
|
} |
|
return 0, nil |
|
} else { |
|
if exitErr, ok := err.(*exec.ExitError); ok { |
|
return exitErr.ExitCode(), nil |
|
} |
|
return -1, err |
|
} |
|
} |
|
|
|
func worker(ctx context.Context, id int, jobs <-chan string, result chan<- string, wg *sync.WaitGroup, processed *uint64, cmdName string, cmdArgsTemplate []string) { |
|
defer wg.Done() |
|
for { |
|
select { |
|
case <-ctx.Done(): |
|
return |
|
case candidate, ok := <-jobs: |
|
if !ok { |
|
return |
|
} |
|
atomic.AddUint64(processed, 1) |
|
var args []string |
|
for _, a := range cmdArgsTemplate { |
|
if a == "{candidate}" { |
|
args = append(args, candidate) |
|
} else { |
|
args = append(args, a) |
|
} |
|
} |
|
exitCode, err := runCmdReturnExitCode(ctx, cmdName, args...) |
|
if err != nil && exitCode == -1 { |
|
continue |
|
} |
|
if exitCode == 0 { |
|
select { |
|
case result <- candidate: |
|
default: |
|
} |
|
return |
|
} |
|
} |
|
} |
|
} |
|
|
|
func main() { |
|
var ( |
|
wordfile = flag.String("wordfile", "", "path to wordlist file") |
|
workers = flag.Int("workers", 100, "number of worker goroutines") |
|
timeout = flag.Duration("timeout", 0, "optional timeout (e.g. 30s). 0 = no timeout") |
|
cmdName = flag.String("cmd", "mycmd", "command to run for each candidate") |
|
) |
|
flag.Parse() |
|
cmdArgsTemplate := flag.Args() |
|
|
|
if *wordfile == "" { |
|
log.Fatalf("usage: %s -wordfile <file> [-workers N] [-timeout 30s] -cmd <command> [args...]\nExample: -cmd /usr/bin/mytool -arg1 -p {candidate}", os.Args[0]) |
|
} |
|
|
|
if *workers <= 0 { |
|
*workers = runtime.NumCPU() |
|
} |
|
f, err := os.Open(*wordfile) |
|
if err != nil { |
|
log.Fatalf("open wordfile: %v", err) |
|
} |
|
defer f.Close() |
|
ctx := context.Background() |
|
if *timeout > 0 { |
|
var cancel context.CancelFunc |
|
ctx, cancel = context.WithTimeout(ctx, *timeout) |
|
defer cancel() |
|
} |
|
ctx, cancel := context.WithCancel(ctx) |
|
defer cancel() |
|
jobs := make(chan string, *workers*2) |
|
var wg sync.WaitGroup |
|
var processed uint64 |
|
result := make(chan string, 1) |
|
for i := 0; i < *workers; i++ { |
|
wg.Add(1) |
|
go worker(ctx, i, jobs, result, &wg, &processed, *cmdName, cmdArgsTemplate) |
|
} |
|
go func() { |
|
scanner := bufio.NewScanner(f) |
|
scanner.Split(bufio.ScanWords) |
|
for scanner.Scan() { |
|
select { |
|
case <-ctx.Done(): |
|
break |
|
default: |
|
jobs <- scanner.Text() |
|
} |
|
} |
|
if err := scanner.Err(); err != nil { |
|
log.Printf("scan error: %v", err) |
|
} |
|
close(jobs) |
|
}() |
|
var found string |
|
select { |
|
case found = <-result: |
|
cancel() |
|
wg.Wait() |
|
case <-ctx.Done(): |
|
wg.Wait() |
|
} |
|
select { |
|
case r := <-result: |
|
found = r |
|
default: |
|
} |
|
if found != "" { |
|
fmt.Printf("FOUND: %s (processed %d candidates)\n", found, atomic.LoadUint64(&processed)) |
|
} else { |
|
fmt.Printf("No success. Processed %d candidates. Context done: %v\n", atomic.LoadUint64(&processed), ctx.Err()) |
|
} |
|
} |