simplify code by using interfaces

gobusterdir/gobusterdir.go

@@ -9,8 +9,11 @@ import (
 	uuid "github.com/satori/go.uuid"
 )
 
-// SetupDir is the setup implementation of gobusterdir
-func SetupDir(g *libgobuster.Gobuster) error {
+// GobusterDir is the main type to implement the interface
+type GobusterDir struct{}
+
+// Setup is the setup implementation of gobusterdir
+func (d GobusterDir) Setup(g *libgobuster.Gobuster) error {
 	_, _, err := g.GetRequest(g.Opts.URL)
 	if err != nil {
 		return fmt.Errorf("unable to connect to %s: %v", g.Opts.URL, err)
@@ -35,8 +38,8 @@ func SetupDir(g *libgobuster.Gobuster) error {
 	return nil
 }
 
-// ProcessDirEntry is the process implementation of gobusterdir
-func ProcessDirEntry(g *libgobuster.Gobuster, word string) ([]libgobuster.Result, error) {
+// Process is the process implementation of gobusterdir
+func (d GobusterDir) Process(g *libgobuster.Gobuster, word string) ([]libgobuster.Result, error) {
 	suffix := ""
 	if g.Opts.UseSlash {
 		suffix = "/"
@@ -77,8 +80,8 @@ func ProcessDirEntry(g *libgobuster.Gobuster, word string) ([]libgobuster.Result
 	return ret, nil
 }
 
-// DirResultToString is the to string implementation of gobusterdir
-func DirResultToString(g *libgobuster.Gobuster, r *libgobuster.Result) (*string, error) {
+// ResultToString is the to string implementation of gobusterdir
+func (d GobusterDir) ResultToString(g *libgobuster.Gobuster, r *libgobuster.Result) (*string, error) {
 	buf := &bytes.Buffer{}
 
 	// Prefix if we're in verbose mode

gobusterdns/gobusterdns.go

@@ -11,8 +11,11 @@ import (
 	uuid "github.com/satori/go.uuid"
 )
 
-// SetupDNS is the setup implementation of gobusterdns
-func SetupDNS(g *libgobuster.Gobuster) error {
+// GobusterDNS is the main type
+type GobusterDNS struct{}
+
+// Setup is the setup implementation of gobusterdns
+func (d GobusterDNS) Setup(g *libgobuster.Gobuster) error {
 	// Resolve a subdomain sthat probably shouldn't exist
 	guid := uuid.Must(uuid.NewV4())
 	wildcardIps, err := net.LookupHost(fmt.Sprintf("%s.%s", guid, g.Opts.URL))
@@ -37,8 +40,8 @@ func SetupDNS(g *libgobuster.Gobuster) error {
 	return nil
 }
 
-// ProcessDNSEntry is the process implementation of gobusterdns
-func ProcessDNSEntry(g *libgobuster.Gobuster, word string) ([]libgobuster.Result, error) {
+// Process is the process implementation of gobusterdns
+func (d GobusterDNS) Process(g *libgobuster.Gobuster, word string) ([]libgobuster.Result, error) {
 	subdomain := fmt.Sprintf("%s.%s", word, g.Opts.URL)
 	ips, err := net.LookupHost(subdomain)
 	var ret []libgobuster.Result
@@ -66,8 +69,8 @@ func ProcessDNSEntry(g *libgobuster.Gobuster, word string) ([]libgobuster.Result
 	return ret, nil
 }
 
-// DNSResultToString is the to string implementation of gobusterdns
-func DNSResultToString(g *libgobuster.Gobuster, r *libgobuster.Result) (*string, error) {
+// ResultToString is the to string implementation of gobusterdns
+func (d GobusterDNS) ResultToString(g *libgobuster.Gobuster, r *libgobuster.Result) (*string, error) {
 	buf := &bytes.Buffer{}
 
 	if r.Status == 404 {

libgobuster/libgobuster.go

@@ -33,16 +33,21 @@ type Gobuster struct {
 	requestsExpected int
 	requestsIssued   int
 	mu               *sync.RWMutex
-	funcResToString  ResultToStringFunc
-	funcProcessor    ProcessFunc
-	funcSetup        SetupFunc
+	plugin           GobusterPlugin
 	IsWildcard       bool
 	resultChan       chan Result
 	errorChan        chan error
 }
 
+// GobusterPlugin is an interface which plugins must implement
+type GobusterPlugin interface {
+	Setup(*Gobuster) error
+	Process(*Gobuster, string) ([]Result, error)
+	ResultToString(*Gobuster, *Result) (*string, error)
+}
+
 // NewGobuster returns a new Gobuster object
-func NewGobuster(c context.Context, opts *Options, setupFunc SetupFunc, processFunc ProcessFunc, resultFunc ResultToStringFunc) (*Gobuster, error) {
+func NewGobuster(c context.Context, opts *Options, plugin GobusterPlugin) (*Gobuster, error) {
 	// validate given options
 	multiErr := opts.validate()
 	if multiErr != nil {
@@ -59,9 +64,7 @@ func NewGobuster(c context.Context, opts *Options, setupFunc SetupFunc, processF
 	}
 	g.http = h
 
-	g.funcSetup = setupFunc
-	g.funcProcessor = processFunc
-	g.funcResToString = resultFunc
+	g.plugin = plugin
 	g.mu = new(sync.RWMutex)
 
 	g.resultChan = make(chan Result)
@@ -122,7 +125,7 @@ func (g *Gobuster) worker(wordChan <-chan string, wg *sync.WaitGroup) {
 				return
 			}
 			// Mode-specific processing
-			res, err := g.funcProcessor(g, word)
+			res, err := g.plugin.Process(g, word)
 			if err != nil {
 				// do not exit and continue
 				g.errorChan <- err
@@ -170,7 +173,7 @@ func (g *Gobuster) getWordlist() (*bufio.Scanner, error) {
 // Start the busting of the website with the given
 // set of settings from the command line.
 func (g *Gobuster) Start() error {
-	if err := g.funcSetup(g); err != nil {
+	if err := g.plugin.Setup(g); err != nil {
 		return err
 	}
 

libgobuster/result.go

@@ -10,7 +10,7 @@ type Result struct {
 
 // ToString converts the Result to it's textual representation
 func (r *Result) ToString(g *Gobuster) (string, error) {
-	s, err := g.funcResToString(g, r)
+	s, err := g.plugin.ResultToString(g, r)
 	if err != nil {
 		return "", err
 	}

main.go

@@ -147,22 +147,15 @@ func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	var funcSetup func(*libgobuster.Gobuster) error
-	var funcProcessor func(*libgobuster.Gobuster, string) ([]libgobuster.Result, error)
-	var funcResToString func(*libgobuster.Gobuster, *libgobuster.Result) (*string, error)
-
+	var plugin libgobuster.GobusterPlugin
 	switch o.Mode {
 	case libgobuster.ModeDir:
-		funcSetup = gobusterdir.SetupDir
-		funcProcessor = gobusterdir.ProcessDirEntry
-		funcResToString = gobusterdir.DirResultToString
+		plugin = gobusterdir.GobusterDir{}
 	case libgobuster.ModeDNS:
-		funcSetup = gobusterdns.SetupDNS
-		funcProcessor = gobusterdns.ProcessDNSEntry
-		funcResToString = gobusterdns.DNSResultToString
+		plugin = gobusterdns.GobusterDNS{}
 	}
 
-	gobuster, err := libgobuster.NewGobuster(ctx, o, funcSetup, funcProcessor, funcResToString)
+	gobuster, err := libgobuster.NewGobuster(ctx, o, plugin)
 	if err != nil {
 		log.Fatalf("[!] %v", err)
 	}