blink/v1/pkg/runner/runner.go

package runner

import (
	"time"

	"git.supernets.org/perp/blink/v1/pkg/dns"
	mdns "github.com/miekg/dns"
	"github.com/panjf2000/ants/v2"
)

// Bruteforce runner
type Runner struct {
	options   *Options         // CLI options
	client    *mdns.Client     // DNS client
	pool      *ants.Pool       // Goroutine pool
	results   chan *dns.Result // Results channel
	wildcards map[string]bool  // Domain wildcards
}

// Return a new Runner
func New(options *Options) *Runner {
	// Check options
	options.Check()

	// Create client
	client := &mdns.Client{
		Net:     "tcp",
		Timeout: time.Second * time.Duration(options.Timeout),
	}

	// Detect UDP
	if options.UDP {
		client.Net = "udp"
	}

	// Create pool
	pool, err := ants.NewPool(options.Threads)
	if err != nil {
		panic(err)
	}

	// Store task count
	tasks := len(options.Wordlist) * len(options.Domains)

	// Create channel
	results := make(chan *dns.Result, tasks)

	// Create wildcards
	wildcards := make(map[string]bool)

	return &Runner{
		options:   options,
		client:    client,
		pool:      pool,
		results:   results,
		wildcards: wildcards,
	}
}

// Detect wildcard
func (r *Runner) Detect() {
	// Go through domains
	for _, domain := range r.options.Domains {
		// Get wildcard status
		r.wildcards[domain] = dns.Wildcard(r.client, r.options.Resolvers, domain)
	}
}

// Submit tasks into pool
func (r *Runner) Submit() {
	// Go through wordlist
	for _, word := range r.options.Wordlist {
		// Go through domains
		for _, domain := range r.options.Domains {
			// Create query
			query := &dns.Query{
				Client:    r.client,
				Resolvers: r.options.Resolvers,
				Domain:    domain,
				Subdomain: word + "." + domain,
				Results:   r.results,
			}

			// Submit query
			r.pool.Submit(func() {
				// Check IPv6
				if !r.options.IPv6 {
					query.A()
				} else {
					query.AAAA()
				}
			})
		}
	}
}

// Receive tasks from pool
func (r *Runner) Receive() {
	// Store task count
	tasks := len(r.options.Wordlist) * len(r.options.Domains)

	// Go through tasks
	for range tasks {
		select {
		case result := <-r.results:
			// Go through domains
			for domain, wildcard := range r.wildcards {
				// Domain found
				if result.Domain == domain {
					// Set wildcard
					result.Wildcard = wildcard
				}
			}

			// Send result
			r.options.OnResult(result)
		}
	}
}

// Start bruteforcing
func (r *Runner) Bruteforce() {
	if r.options.Wildcard {
		r.Detect()
	}

	go r.Submit()
	r.Receive()
}