Skip to content

Commit

Permalink
Refactored afew places
Browse files Browse the repository at this point in the history
  • Loading branch information
kenjoe41 committed Sep 13, 2024
1 parent 6e4c086 commit e5caad9
Show file tree
Hide file tree
Showing 9 changed files with 241 additions and 129 deletions.
55 changes: 43 additions & 12 deletions Makefile
Original file line number Diff line number Diff line change
@@ -1,20 +1,51 @@
tests:
go test . ./... -v -race -covermode atomic -coverprofile coverage.out && go tool cover -html coverage.out -o coverage.html
# Variables for coverage and profiling files
COVERAGE_FILE := coverage.out
COVERAGE_HTML := coverage.html
CPU_PROFILE := cpu.prof
MEM_PROFILE := mem.prof

tests_without_race:
go test . ./... -v -covermode atomic -coverprofile coverage.out && go tool cover -html coverage.out -o coverage.html
# Default target runs tests with race conditions
.PHONY: all
all: test

# Test with race detection, coverage report generation
.PHONY: test
test:
go test ./... -v -race -covermode=atomic -coverprofile=$(COVERAGE_FILE)
go tool cover -html=$(COVERAGE_FILE) -o $(COVERAGE_HTML)

# Test without race detection, coverage report generation
.PHONY: test-no-race
test-no-race:
go test ./... -v -covermode=atomic -coverprofile=$(COVERAGE_FILE)
go tool cover -html=$(COVERAGE_FILE) -o $(COVERAGE_HTML)

# Format the code
.PHONY: format
format:
go fmt . ./...
go fmt ./...

# Run benchmarks with memory allocation statistics
.PHONY: bench
bench:
go test . ./... -bench . -benchmem -cpu 1
go test ./... -bench . -benchmem -cpu=1

# Generate CPU and memory profiles while running benchmarks
.PHONY: profile-bench
profile-bench:
go test ./... -bench . -cpuprofile=$(CPU_PROFILE) -memprofile=$(MEM_PROFILE) -cpu=1

report_bench:
go test . ./... -cpuprofile cpu.prof -memprofile mem.prof -bench . -cpu 1
# Generate CPU profiling report
.PHONY: cpu-report
cpu-report:
go tool pprof $(CPU_PROFILE)

cpu_report:
go tool pprof cpu.prof
# Generate memory profiling report
.PHONY: mem-report
mem-report:
go tool pprof $(MEM_PROFILE)

mem_report:
go tool pprof mem.prof
# Clean up profiling and coverage files
.PHONY: clean
clean:
rm -f $(COVERAGE_FILE) $(COVERAGE_HTML) $(CPU_PROFILE) $(MEM_PROFILE)
158 changes: 74 additions & 84 deletions cli/cli.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@ package cli

import (
"bufio"
"context"
"errors"
"flag"
"fmt"
"log"
"os"
Expand All @@ -16,122 +16,112 @@ import (
"github.com/kenjoe41/goSubsWordlist/output"
)

// Cli accepts a list of URLs, one URL per line, from stdin and generates a wordlist from all subdomains found in the list.
func Cli(includeRoot, silent bool) error {
// Print Header text
if !silent {
output.Beautify()
}

// This is a CPU-bound task, increasing the threads beyond what's available will just make it slow so removed the flag option.
concurrency := runtime.NumCPU()

// This is divided up in the subroutine for loop, so a value below 2 is BS.
if concurrency < 2 {
concurrency = 2
} else {
// We have 2 channels to share the concurrency with, let's reassure them that they'll have equal share.
concurrency *= 2
}

// Create channels to use
// Channels and WaitGroups
domains := make(chan string)
subdomains := make(chan string)
output := make(chan string)

// Domain Input worker
var domainsWG sync.WaitGroup
for i := 0; i < concurrency/2; i++ {
domainsWG.Add(1)
words := make(chan string)

go func() {
extract, err := fasttld.New(fasttld.SuffixListParams{})
if err != nil {
log.Fatal(err) // unlikely
}
for domain := range domains {
if domain == "" {
// Log something but continue to next domain if available
// log.Printf("Failed to get domain from: %s", domain)
continue
}
subdomain := ezutils.ExtractSubdomain(domain, includeRoot, extract)

if subdomain == "" {
// Log something but continue to next domain if available
// log.Printf("Failed to get subdomain for domain: %s", domain)
continue
}
subdomains <- subdomain
}
domainsWG.Done()
}()
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()

var subdomainsWG sync.WaitGroup
var wg sync.WaitGroup

// Domain input workers
wg.Add(concurrency / 2)
for i := 0; i < concurrency/2; i++ {
subdomainsWG.Add(1)

go func() {
for inSubdomains := range subdomains {
// Split the subdomain into separate words by the '.' char.
// Returns slice of words.
subWords := strings.Split(inSubdomains, ".")

// Print to console for now
for _, subword := range subWords {
output <- subword
}
}
subdomainsWG.Done()
}()
go processDomains(ctx, &wg, domains, subdomains, includeRoot)
}

// Close subdomains channel when done reading from domains chan.
go func() {
domainsWG.Wait()
close(subdomains)
}()
// Subdomain processing workers
wg.Add(concurrency / 2)
for i := 0; i < concurrency/2; i++ {
go processSubdomains(ctx, &wg, subdomains, words)
}

var outputWG sync.WaitGroup
outputWG.Add(1)
// Output processor
go func() {
for word := range output {
for word := range words {
fmt.Println(word)
}
outputWG.Done()
}()

// Close the Output Chan after subdomain worker is done.
go func() {
subdomainsWG.Wait()
close(output)
}()
// Read input from stdin
if err := readStdin(domains); err != nil {
return err
}
close(domains)

// Check for stdin input
stat, _ := os.Stdin.Stat()
if (stat.Mode() & os.ModeCharDevice) != 0 {
flag.Usage()
return errors.New("No domains or urls detected. Hint: cat domains.txt | goSubsWordlist")
// Wait for workers to complete
wg.Wait()
close(subdomains)
close(words)

return nil
}

func processDomains(ctx context.Context, wg *sync.WaitGroup, domains <-chan string, subdomains chan<- string, includeRoot bool) {
defer wg.Done()
extract, err := fasttld.New(fasttld.SuffixListParams{})
if err != nil {
log.Fatal(err) // unlikely
}
for {
select {
case <-ctx.Done():
return
case domain, ok := <-domains:
if !ok {
return
}
if subdomain := ezutils.ExtractSubdomain(domain, includeRoot, extract); subdomain != "" {
subdomains <- subdomain
}
}
}
}

sc := bufio.NewScanner(os.Stdin)
func processSubdomains(ctx context.Context, wg *sync.WaitGroup, subdomains <-chan string, words chan<- string) {
defer wg.Done()
for {
select {
case <-ctx.Done():
return
case subdomain, ok := <-subdomains:
if !ok {
return
}
subWords := strings.Split(subdomain, ".")
for _, word := range subWords {
words <- word
}
}
}
}

for sc.Scan() {
domains <- sc.Text()
func readStdin(domains chan<- string) error {
stat, _ := os.Stdin.Stat()
if (stat.Mode() & os.ModeCharDevice) != 0 {
return errors.New("No domains or URLs detected. Hint: cat domains.txt | goSubsWordlist")
}

// Close domains chan
close(domains)
scanner := bufio.NewScanner(os.Stdin)
for scanner.Scan() {
domains <- scanner.Text()
}

// check there were no errors reading stdin (unlikely)
if err := sc.Err(); err != nil {
if err := scanner.Err(); err != nil {
return err
}

// Wait until the output waitgroup is done
outputWG.Wait()

return nil
}
65 changes: 56 additions & 9 deletions cli/cli_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,21 +5,70 @@ import (
"testing"
)

func TestCLI(t *testing.T) {
func TestCLI_NoStdin(t *testing.T) {
// Test when there is no input from Stdin
if err := Cli(false, false); err == nil {
t.Error("Expected error for no Stdin input")
}
}

func TestCLI_WithStdin(t *testing.T) {
// Test with simulated Stdin input
userInput := "example.com\nsub.example.com\nanother.sub.example.com\n\n"

funcDefer, err := mockStdin(t, userInput)
if err != nil {
t.Fatalf("Error mocking Stdin: %v", err)
}
defer funcDefer()

if err := Cli(false, false); err != nil {
t.Errorf("Unexpected error: %v", err)
}
}

func TestCLI_WithIncludeRoot(t *testing.T) {
// Test including the root domain in the wordlist
userInput := "example.com\nsub.example.com\n\n"

funcDefer, err := mockStdin(t, userInput)
if err != nil {
t.Fatalf("Error mocking Stdin: %v", err)
}
defer funcDefer()

if err := Cli(true, false); err != nil {
t.Errorf("Unexpected error: %v", err)
}
}

func TestCLI_WithSilentFlag(t *testing.T) {
// Test with the silent flag enabled (no banner output)
userInput := "example.com\nsub.example.com\n\n"

funcDefer, err := mockStdin(t, userInput)
if err != nil {
t.Errorf("%q", err)
t.Fatalf("Error mocking Stdin: %v", err)
}
defer funcDefer()

if err := Cli(false, true); err != nil {
t.Errorf("Unexpected error: %v", err)
}
}

// Test with invalid input
func TestCLI_InvalidInput(t *testing.T) {
userInput := "invalid-url"

funcDefer, err := mockStdin(t, userInput)
if err != nil {
t.Fatalf("Error mocking Stdin: %v", err)
}
defer funcDefer()

if err := Cli(false, false); err != nil {
t.Errorf("%q", err)
t.Errorf("Unexpected error: %v", err)
}
}

Expand All @@ -30,27 +79,25 @@ func mockStdin(t *testing.T, dummyInput string) (funcDefer func(), err error) {

oldOsStdin := os.Stdin
tmpfile, err := os.CreateTemp(t.TempDir(), t.Name())

if err != nil {
return nil, err
}

content := []byte(dummyInput)

if _, err := tmpfile.Write(content); err != nil {
if _, err := tmpfile.Write([]byte(dummyInput)); err != nil {
return nil, err
}

if _, err := tmpfile.Seek(0, 0); err != nil {
return nil, err
}

// Set stdin to the temp file
// Set os.Stdin to the temp file
os.Stdin = tmpfile

return func() {
// clean up
// Clean up
os.Stdin = oldOsStdin
tmpfile.Close()
os.Remove(tmpfile.Name())
}, nil
}
9 changes: 3 additions & 6 deletions ezutils/ezutils.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,11 @@ import (
"github.com/elliotwutingfeng/go-fasttld"
)

// ExtractSubdomain extracts the subdomain from a given url.
// If includeRootPtr is true, the second-level domain will be included
// ExtractSubdomain extracts the subdomain from a given URL.
// If includeRootPtr is true, the second-level domain will be included in the result.
func ExtractSubdomain(url string, includeRootPtr bool, extract *fasttld.FastTLD) string {
result, err := extract.Extract(fasttld.URLParams{URL: url})
if err != nil {
return ""
}
if result.HostType != fasttld.HostName {
if err != nil || result.HostType != fasttld.HostName {
return ""
}
if len(result.SubDomain) > 0 {
Expand Down
Loading

0 comments on commit e5caad9

Please sign in to comment.