diff --git a/pkg/massdns/massdns.go b/pkg/massdns/massdns.go index 4c05b8b..2b6bb8e 100644 --- a/pkg/massdns/massdns.go +++ b/pkg/massdns/massdns.go @@ -46,6 +46,10 @@ type Options struct { WildcardOutputFile string // MassDnsCmd supports massdns flags MassDnsCmd string + // KeepStderr controls whether to capture and store massdns stderr output + KeepStderr bool + // BatchSize controls the number of lines per chunk for incremental processing + BatchSize int OnResult func(*retryabledns.DNSData) } diff --git a/pkg/massdns/process.go b/pkg/massdns/process.go index 0edb1ae..ef9a9ea 100644 --- a/pkg/massdns/process.go +++ b/pkg/massdns/process.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "os" "os/exec" "strconv" @@ -19,9 +20,7 @@ import ( "github.com/projectdiscovery/shuffledns/pkg/wildcards" "github.com/projectdiscovery/utils/batcher" fileutil "github.com/projectdiscovery/utils/file" - folderutil "github.com/projectdiscovery/utils/folder" ioutil "github.com/projectdiscovery/utils/io" - stringsutil "github.com/projectdiscovery/utils/strings" "github.com/remeh/sizedwaitgroup" "github.com/weppos/publicsuffix-go/publicsuffix" ) @@ -30,32 +29,53 @@ import ( func (instance *Instance) RunWithContext(ctx context.Context) (stdout, stderr string, took time.Duration, err error) { start := time.Now() + // Create temporary file for massdns output stdoutFile, err := os.CreateTemp(instance.options.TempDir, "massdns-stdout-") if err != nil { - return "", "", 0, fmt.Errorf("could not create temp file for massdns stdout: %w", err) + return "", "", 0, fmt.Errorf("could not create temp file for massdns output: %w", err) } defer func() { _ = stdoutFile.Close() }() - stderrFile, err := os.CreateTemp(instance.options.TempDir, "massdns-stderr-") - if err != nil { - return "", "", 0, fmt.Errorf("could not create temp file for massdns stdout: %w", err) + // Handle stderr based on KeepStderr option + var stderrFile *os.File + if instance.options.KeepStderr { + stderrFile, err = os.CreateTemp(instance.options.TempDir, "massdns-stderr-") + if err != nil { + return "", "", 0, fmt.Errorf("could not create temp file for massdns stderr: %w", err) + } + defer func() { + _ = stderrFile.Close() + }() } - defer func() { - _ = stderrFile.Close() - }() // Run the command on a temp file and wait for the output args := []string{"-r", instance.options.ResolversFile, "-o", "Snl", "--retry", "REFUSED", "--retry", "SERVFAIL", "-t", "A", instance.options.InputFile, "-s", strconv.Itoa(instance.options.Threads)} if instance.options.MassDnsCmd != "" { args = append(args, strings.Split(instance.options.MassDnsCmd, " ")...) } + cmd := exec.CommandContext(ctx, instance.options.MassdnsPath, args...) cmd.Stdout = stdoutFile - cmd.Stderr = stderrFile + + // Set stderr based on KeepStderr option + if instance.options.KeepStderr { + cmd.Stderr = stderrFile + } else { + // Discard stderr by sending it to /dev/null + cmd.Stderr = nil + } + err = cmd.Run() - return stdoutFile.Name(), stderrFile.Name(), time.Since(start), err + + // Return stderr filename only if it was captured + stderrFilename := "" + if instance.options.KeepStderr { + stderrFilename = stderrFile.Name() + } + + return stdoutFile.Name(), stderrFilename, time.Since(start), err } func (instance *Instance) Run(ctx context.Context) error { @@ -81,38 +101,11 @@ func (instance *Instance) Run(ctx context.Context) error { } defer shstore.Close() - // Set the correct target file - tmpDir := instance.options.TempDir - // Check if we need to run massdns if instance.options.MassdnsRaw == "" { - if len(instance.options.Domains) > 0 { - gologger.Info().Msgf("Executing massdns on %s\n", strings.Join(instance.options.Domains, ", ")) - } else { - gologger.Info().Msgf("Executing massdns\n") - } - - // Create a temporary file for the massdns output - gologger.Info().Msgf("using massdns output directory: %s\n", tmpDir) - stdoutFile, stderrFile, took, err := instance.RunWithContext(ctx) - gologger.Info().Msgf("massdns output file: %s\n", stdoutFile) - gologger.Info().Msgf("massdns error file: %s\n", stderrFile) - if err != nil { - return fmt.Errorf("could not execute massdns: %s", err) - } - - gologger.Info().Msgf("Massdns execution took %s\n", took) - - gologger.Info().Msgf("Started parsing massdns output\n") - - now := time.Now() - - err = instance.parseMassDNSOutputDir(tmpDir, shstore) - if err != nil { - return fmt.Errorf("could not parse massdns output: %w", err) - } - - gologger.Info().Msgf("Massdns output parsing completed in %s\n", time.Since(now)) + // This case is now handled by the streaming methods in the runner + // The Run method is only called for raw massdns output processing + return errors.New("streaming processing should be used for new massdns runs") } else { // parse the input file gologger.Info().Msgf("Started parsing massdns input\n") now := time.Now() @@ -156,6 +149,59 @@ func (instance *Instance) Run(ctx context.Context) error { return nil } +// runChunk runs massdns on a specific chunk file +func (instance *Instance) runChunk(ctx context.Context, chunkFile string) (stdout, stderr string, took time.Duration, err error) { + start := time.Now() + + // Create temporary file for massdns output + stdoutFile, err := os.CreateTemp(instance.options.TempDir, "massdns-chunk-stdout-") + if err != nil { + return "", "", 0, fmt.Errorf("could not create temp file for massdns output: %w", err) + } + defer func() { + _ = stdoutFile.Close() + }() + + // Handle stderr based on KeepStderr option + var stderrFile *os.File + if instance.options.KeepStderr { + stderrFile, err = os.CreateTemp(instance.options.TempDir, "massdns-chunk-stderr-") + if err != nil { + return "", "", 0, fmt.Errorf("could not create temp file for massdns stderr: %w", err) + } + defer func() { + _ = stderrFile.Close() + }() + } + + // Run the command on the chunk file + args := []string{"-r", instance.options.ResolversFile, "-o", "Snl", "--retry", "REFUSED", "--retry", "SERVFAIL", "-t", "A", chunkFile, "-s", strconv.Itoa(instance.options.Threads)} + if instance.options.MassDnsCmd != "" { + args = append(args, strings.Split(instance.options.MassDnsCmd, " ")...) + } + + cmd := exec.CommandContext(ctx, instance.options.MassdnsPath, args...) + cmd.Stdout = stdoutFile + + // Set stderr based on KeepStderr option + if instance.options.KeepStderr { + cmd.Stderr = stderrFile + } else { + // Discard stderr by sending it to /dev/null + cmd.Stderr = nil + } + + err = cmd.Run() + + // Return stderr filename only if it was captured + stderrFilename := "" + if instance.options.KeepStderr { + stderrFilename = stderrFile.Name() + } + + return stdoutFile.Name(), stderrFilename, time.Since(start), err +} + type item struct { ip string domain string @@ -205,26 +251,6 @@ func (instance *Instance) parseMassDNSOutputFile(tmpFile string, store *store.St return nil } -func (instance *Instance) parseMassDNSOutputDir(tmpDir string, store *store.Store) error { - tmpFiles, err := folderutil.GetFiles(tmpDir) - if err != nil { - return fmt.Errorf("could not open massdns output directory: %w", err) - } - - for _, tmpFile := range tmpFiles { - // just process stdout files - if !stringsutil.ContainsAnyI(tmpFile, "stdout") { - continue - } - err = instance.parseMassDNSOutputFile(tmpFile, store) - if err != nil { - return fmt.Errorf("could not parse massdns output: %w", err) - } - } - - return nil -} - func (instance *Instance) autoExtractRootDomains(store *store.Store) error { candidateRootDomains := make(map[string]struct{}) store.Iterate(func(ip string, hostnames []string, counter int) { @@ -436,3 +462,293 @@ func (instance *Instance) writeOutput(store *store.Store) error { } return nil } + +// ProcessDomainStreaming processes domain bruteforce using streaming with batcher +func (instance *Instance) ProcessDomainStreaming(ctx context.Context, wordlistFile *os.File) error { + // Create a store for storing ip metadata + shstore, err := store.New(instance.options.TempDir) + if err != nil { + return fmt.Errorf("could not create store: %w", err) + } + defer shstore.Close() + + // Create batcher for streaming permutations + chunkNum := 0 + permutationCount := 0 + + bulkWriter := batcher.New[string]( + batcher.WithMaxCapacity[string](instance.options.BatchSize), + batcher.WithFlushInterval[string](10*time.Second), + batcher.WithFlushCallback[string](func(permutations []string) { + chunkNum++ + if len(permutations) == 0 { + return + } + + gologger.Info().Msgf("Processing chunk %d (%d permutations, total: %d)\n", + chunkNum, len(permutations), permutationCount) + + // Create temporary chunk file + chunkFile, err := os.CreateTemp(instance.options.TempDir, fmt.Sprintf("chunk-%d-", chunkNum)) + if err != nil { + gologger.Error().Msgf("Could not create chunk file: %s\n", err) + return + } + + // Write permutations to chunk file + writer := bufio.NewWriter(chunkFile) + for _, permutation := range permutations { + _, err := writer.WriteString(permutation + "\n") + if err != nil { + gologger.Error().Msgf("Could not write to chunk file: %s\n", err) + _ = chunkFile.Close() + _ = os.Remove(chunkFile.Name()) + return + } + } + _ = writer.Flush() + _ = chunkFile.Close() + + // Run massdns on this chunk + chunkStart := time.Now() + stdoutFile, stderrFile, took, err := instance.runChunk(ctx, chunkFile.Name()) + if err != nil { + gologger.Error().Msgf("Could not execute massdns on chunk %d: %s\n", chunkNum, err) + _ = os.Remove(chunkFile.Name()) + return + } + + gologger.Info().Msgf("Chunk %d massdns execution took %s\n", chunkNum, took) + + // Parse the chunk output immediately + parseStart := time.Now() + err = instance.parseMassDNSOutputFile(stdoutFile, shstore) + if err != nil { + gologger.Error().Msgf("Could not parse massdns output for chunk %d: %s\n", chunkNum, err) + _ = os.Remove(chunkFile.Name()) + _ = os.Remove(stdoutFile) + if stderrFile != "" { + _ = os.Remove(stderrFile) + } + return + } + + gologger.Info().Msgf("Chunk %d parsing completed in %s\n", chunkNum, time.Since(parseStart)) + + // Clean up chunk files immediately + _ = os.Remove(chunkFile.Name()) + _ = os.Remove(stdoutFile) + if stderrFile != "" { + _ = os.Remove(stderrFile) + } + + gologger.Info().Msgf("Chunk %d completed in %s\n", chunkNum, time.Since(chunkStart)) + }), + ) + + bulkWriter.Run() + + // Read wordlist and generate permutations on-the-fly + scanner := bufio.NewScanner(wordlistFile) + for scanner.Scan() { + // RFC4343 - case insensitive domain + text := strings.ToLower(scanner.Text()) + if text == "" { + continue + } + + // Generate permutations for each domain + for _, domain := range instance.options.Domains { + permutation := text + "." + domain + bulkWriter.Append(permutation) + permutationCount++ + } + } + + // Stop the batcher and wait for completion + bulkWriter.Stop() + bulkWriter.WaitDone() + + if err := scanner.Err(); err != nil { + return fmt.Errorf("error reading wordlist: %w", err) + } + + gologger.Info().Msgf("Total permutations generated: %d\n", permutationCount) + + // Perform post-processing steps + if instance.options.AutoExtractRootDomains { + gologger.Info().Msgf("Started extracting root domains\n") + now := time.Now() + err = instance.autoExtractRootDomains(shstore) + if err != nil { + return fmt.Errorf("could not extract root domains: %w", err) + } + gologger.Info().Msgf("Root domain extraction completed in %s\n", time.Since(now)) + } + + // Perform wildcard filtering only if domain name has been specified + if len(instance.options.Domains) > 0 { + gologger.Info().Msgf("Started removing wildcards records\n") + now := time.Now() + err = instance.filterWildcards(shstore) + if err != nil { + return fmt.Errorf("could not filter wildcards: %w", err) + } + gologger.Info().Msgf("Wildcard removal completed in %s\n", time.Since(now)) + } + + gologger.Info().Msgf("Finished enumeration, started writing output\n") + + // Write the final elaborated list out + now := time.Now() + err = instance.writeOutput(shstore) + if err != nil { + return fmt.Errorf("could not write output: %w", err) + } + gologger.Info().Msgf("Output written in %s\n", time.Since(now)) + + return nil +} + +// ProcessSubdomainsStreaming processes subdomain list using streaming with batcher +func (instance *Instance) ProcessSubdomainsStreaming(ctx context.Context, subdomainReader io.Reader) error { + // Create a store for storing ip metadata + shstore, err := store.New(instance.options.TempDir) + if err != nil { + return fmt.Errorf("could not create store: %w", err) + } + defer shstore.Close() + + // Create batcher for streaming subdomains + chunkNum := 0 + subdomainCount := 0 + + bulkWriter := batcher.New[string]( + batcher.WithMaxCapacity[string](instance.options.BatchSize), + batcher.WithFlushInterval[string](10*time.Second), + batcher.WithFlushCallback[string](func(subdomains []string) { + chunkNum++ + if len(subdomains) == 0 { + return + } + + gologger.Info().Msgf("Processing chunk %d (%d subdomains, total: %d)\n", + chunkNum, len(subdomains), subdomainCount) + + // Create temporary chunk file + chunkFile, err := os.CreateTemp(instance.options.TempDir, fmt.Sprintf("chunk-%d-", chunkNum)) + if err != nil { + gologger.Error().Msgf("Could not create chunk file: %s\n", err) + return + } + + // Write subdomains to chunk file + writer := bufio.NewWriter(chunkFile) + for _, subdomain := range subdomains { + _, err := writer.WriteString(subdomain + "\n") + if err != nil { + gologger.Error().Msgf("Could not write to chunk file: %s\n", err) + _ = chunkFile.Close() + _ = os.Remove(chunkFile.Name()) + return + } + } + _ = writer.Flush() + _ = chunkFile.Close() + + // Run massdns on this chunk + chunkStart := time.Now() + stdoutFile, stderrFile, took, err := instance.runChunk(ctx, chunkFile.Name()) + if err != nil { + gologger.Error().Msgf("Could not execute massdns on chunk %d: %s\n", chunkNum, err) + _ = os.Remove(chunkFile.Name()) + return + } + + gologger.Info().Msgf("Chunk %d massdns execution took %s\n", chunkNum, took) + + // Parse the chunk output immediately + parseStart := time.Now() + err = instance.parseMassDNSOutputFile(stdoutFile, shstore) + if err != nil { + gologger.Error().Msgf("Could not parse massdns output for chunk %d: %s\n", chunkNum, err) + _ = os.Remove(chunkFile.Name()) + _ = os.Remove(stdoutFile) + if stderrFile != "" { + _ = os.Remove(stderrFile) + } + return + } + + gologger.Info().Msgf("Chunk %d parsing completed in %s\n", chunkNum, time.Since(parseStart)) + + // Clean up chunk files immediately + _ = os.Remove(chunkFile.Name()) + _ = os.Remove(stdoutFile) + if stderrFile != "" { + _ = os.Remove(stderrFile) + } + + gologger.Info().Msgf("Chunk %d completed in %s\n", chunkNum, time.Since(chunkStart)) + }), + ) + + bulkWriter.Run() + + // Read subdomains and stream them to batcher + scanner := bufio.NewScanner(subdomainReader) + for scanner.Scan() { + // RFC4343 - case insensitive domain + subdomain := strings.ToLower(strings.TrimSpace(scanner.Text())) + if subdomain == "" { + continue + } + + bulkWriter.Append(subdomain) + subdomainCount++ + } + + // Stop the batcher and wait for completion + bulkWriter.Stop() + bulkWriter.WaitDone() + + if err := scanner.Err(); err != nil { + return fmt.Errorf("error reading subdomains: %w", err) + } + + gologger.Info().Msgf("Total subdomains processed: %d\n", subdomainCount) + + // Perform post-processing steps + if instance.options.AutoExtractRootDomains { + gologger.Info().Msgf("Started extracting root domains\n") + now := time.Now() + err = instance.autoExtractRootDomains(shstore) + if err != nil { + return fmt.Errorf("could not extract root domains: %w", err) + } + gologger.Info().Msgf("Root domain extraction completed in %s\n", time.Since(now)) + } + + // Perform wildcard filtering only if domain name has been specified + if len(instance.options.Domains) > 0 { + gologger.Info().Msgf("Started removing wildcards records\n") + now := time.Now() + err = instance.filterWildcards(shstore) + if err != nil { + return fmt.Errorf("could not filter wildcards: %w", err) + } + gologger.Info().Msgf("Wildcard removal completed in %s\n", time.Since(now)) + } + + gologger.Info().Msgf("Finished enumeration, started writing output\n") + + // Write the final elaborated list out + now := time.Now() + err = instance.writeOutput(shstore) + if err != nil { + return fmt.Errorf("could not write output: %w", err) + } + gologger.Info().Msgf("Output written in %s\n", time.Since(now)) + + return nil +} diff --git a/pkg/runner/options.go b/pkg/runner/options.go index e7a78df..6a18bdb 100644 --- a/pkg/runner/options.go +++ b/pkg/runner/options.go @@ -9,6 +9,11 @@ import ( updateutils "github.com/projectdiscovery/utils/update" ) +const ( + // DefaultBatchSize is the default number of lines per chunk for incremental processing + DefaultBatchSize = 500000 +) + // Options contains the configuration options for tuning // the active dns resolving process. type Options struct { @@ -35,6 +40,8 @@ type Options struct { MassDnsCmd string // Supports massdns flags(example -i) DisableUpdateCheck bool // DisableUpdateCheck disable automatic update check Mode string + KeepStderr bool // KeepStderr controls whether to capture and store massdns stderr output + BatchSize int // BatchSize controls the number of lines per chunk for incremental processing OnResult func(*retryabledns.DNSData) } @@ -43,6 +50,7 @@ var DefaultOptions = Options{ Threads: 10000, Retries: 5, WildcardThreads: 250, + BatchSize: DefaultBatchSize, // Default batch size for incremental processing } // ParseOptions parses the command line flags provided by a user @@ -88,6 +96,8 @@ func ParseOptions() *Options { flagSet.IntVar(&options.Retries, "retries", 5, "Number of retries for dns enumeration"), flagSet.BoolVarP(&options.StrictWildcard, "strict-wildcard", "sw", false, "Perform wildcard check on all found subdomains"), flagSet.IntVar(&options.WildcardThreads, "wt", 250, "Number of concurrent wildcard checks"), + flagSet.BoolVar(&options.KeepStderr, "retain-stderr", false, "Capture and store massdns stderr output (default: discard)"), + flagSet.IntVar(&options.BatchSize, "batch-size", DefaultBatchSize, "Number of lines per chunk for incremental processing"), ) flagSet.CreateGroup("debug", "Debug", diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go index ac70d80..6389e39 100644 --- a/pkg/runner/runner.go +++ b/pkg/runner/runner.go @@ -1,20 +1,14 @@ package runner import ( - "bufio" "context" "errors" - "io" "os" "os/exec" - "path/filepath" - "strings" - "time" "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/shuffledns/pkg/massdns" fileutil "github.com/projectdiscovery/utils/file" - "github.com/rs/xid" ) // Runner is a client for running the enumeration process. @@ -81,9 +75,9 @@ func (r *Runner) findBinary() string { // RunEnumeration sets up the input layer for giving input to massdns // binary and runs the actual enumeration func (r *Runner) RunEnumeration() { - // Handle only wildcard filtering + // Handle only wildcard filtering on existing massdns output if r.options.MassdnsRaw != "" { - r.processSubdomains() + r.processExistingOutput() return } @@ -102,72 +96,123 @@ func (r *Runner) RunEnumeration() { // processDomain processes the bruteforce for a domain using a wordlist func (r *Runner) processDomain() { - resolveFile := filepath.Join(r.tempDir, xid.New().String()) - file, err := os.Create(resolveFile) - if err != nil { - gologger.Error().Msgf("Could not create bruteforce list (%s): %s\n", r.tempDir, err) - return - } - writer := bufio.NewWriter(file) - // Read the input wordlist for bruteforce generation inputFile, err := os.Open(r.options.Wordlist) if err != nil { gologger.Error().Msgf("Could not read bruteforce wordlist (%s): %s\n", r.options.Wordlist, err) - _ = file.Close() return } + defer func() { + _ = inputFile.Close() + }() - gologger.Info().Msgf("Started generating bruteforce permutation\n") + gologger.Info().Msgf("Started generating bruteforce permutation with streaming processing\n") - now := time.Now() - // Create permutation for domain with wordlist - scanner := bufio.NewScanner(inputFile) - for scanner.Scan() { - // RFC4343 - case insensitive domain - text := strings.ToLower(scanner.Text()) - if text == "" { - continue - } - for _, domain := range r.options.Domains { - _, _ = writer.WriteString(text + "." + domain + "\n") - } + // Create massdns instance for processing chunks + massdns, err := massdns.New(massdns.Options{ + Domains: r.options.Domains, + AutoExtractRootDomains: r.options.AutoExtractRootDomains, + Retries: r.options.Retries, + MassdnsPath: r.options.MassdnsPath, + Threads: r.options.Threads, + WildcardsThreads: r.options.WildcardThreads, + ResolversFile: r.options.ResolversFile, + TrustedResolvers: r.options.TrustedResolvers, + TempDir: r.tempDir, + OutputFile: r.options.Output, + Json: r.options.Json, + MassdnsRaw: r.options.MassdnsRaw, + StrictWildcard: r.options.StrictWildcard, + WildcardOutputFile: r.options.WildcardOutputFile, + MassDnsCmd: r.options.MassDnsCmd, + KeepStderr: r.options.KeepStderr, + BatchSize: r.options.BatchSize, + OnResult: r.options.OnResult, + }) + if err != nil { + gologger.Error().Msgf("Could not create massdns client: %s\n", err) + return } - _ = writer.Flush() - _ = inputFile.Close() - _ = file.Close() - gologger.Info().Msgf("Generating permutations took %s at %s\n", time.Since(now), resolveFile) + // Use streaming processing with batcher + err = massdns.ProcessDomainStreaming(context.Background(), inputFile) + if err != nil { + gologger.Error().Msgf("Could not process domain with streaming: %s\n", err) + return + } - // Run the actual massdns enumeration process - r.runMassdns(resolveFile) + if r.options.WildcardOutputFile != "" { + _ = massdns.DumpWildcardsToFile(r.options.WildcardOutputFile) + } + + gologger.Info().Msgf("Finished resolving.\n") } // processSubdomain processes the resolving for a list of subdomains func (r *Runner) processSubdomains() { - var resolveFile string + // Create massdns instance for processing chunks + massdns, err := massdns.New(massdns.Options{ + Domains: r.options.Domains, + AutoExtractRootDomains: r.options.AutoExtractRootDomains, + Retries: r.options.Retries, + MassdnsPath: r.options.MassdnsPath, + Threads: r.options.Threads, + WildcardsThreads: r.options.WildcardThreads, + ResolversFile: r.options.ResolversFile, + TrustedResolvers: r.options.TrustedResolvers, + TempDir: r.tempDir, + OutputFile: r.options.Output, + Json: r.options.Json, + MassdnsRaw: r.options.MassdnsRaw, + StrictWildcard: r.options.StrictWildcard, + WildcardOutputFile: r.options.WildcardOutputFile, + MassDnsCmd: r.options.MassDnsCmd, + KeepStderr: r.options.KeepStderr, + BatchSize: r.options.BatchSize, + OnResult: r.options.OnResult, + }) + if err != nil { + gologger.Error().Msgf("Could not create massdns client: %s\n", err) + return + } - // If there is stdin, write the resolution list to the file + // Handle stdin or file input if fileutil.HasStdin() && r.options.SubdomainsList == "" { - file, err := os.CreateTemp(r.tempDir, "massdns-stdin-") + // Use streaming processing for stdin + gologger.Info().Msgf("Processing subdomains from stdin with streaming\n") + err = massdns.ProcessSubdomainsStreaming(context.Background(), os.Stdin) if err != nil { - gologger.Error().Msgf("Could not create resolution list (%s): %s\n", r.tempDir, err) + gologger.Error().Msgf("Could not process subdomains with streaming: %s\n", err) return } - _, _ = io.Copy(file, os.Stdin) - _ = file.Close() - resolveFile = file.Name() } else { - // Use the file if user has provided one - resolveFile = r.options.SubdomainsList + // Use streaming processing for file + subdomainFile, err := os.Open(r.options.SubdomainsList) + if err != nil { + gologger.Error().Msgf("Could not open subdomain list (%s): %s\n", r.options.SubdomainsList, err) + return + } + defer func() { + _ = subdomainFile.Close() + }() + + gologger.Info().Msgf("Processing subdomains from file with streaming\n") + err = massdns.ProcessSubdomainsStreaming(context.Background(), subdomainFile) + if err != nil { + gologger.Error().Msgf("Could not process subdomains with streaming from file: %s\n", err) + return + } } - // Run the actual massdns enumeration process - r.runMassdns(resolveFile) + if r.options.WildcardOutputFile != "" { + _ = massdns.DumpWildcardsToFile(r.options.WildcardOutputFile) + } + + gologger.Info().Msgf("Finished resolving.\n") } -// runMassdns runs the massdns tool on the list of inputs -func (r *Runner) runMassdns(inputFile string) { +// processExistingOutput processes existing massdns output for wildcard filtering +func (r *Runner) processExistingOutput() { massdns, err := massdns.New(massdns.Options{ Domains: r.options.Domains, AutoExtractRootDomains: r.options.AutoExtractRootDomains, @@ -175,7 +220,6 @@ func (r *Runner) runMassdns(inputFile string) { MassdnsPath: r.options.MassdnsPath, Threads: r.options.Threads, WildcardsThreads: r.options.WildcardThreads, - InputFile: inputFile, ResolversFile: r.options.ResolversFile, TrustedResolvers: r.options.TrustedResolvers, TempDir: r.tempDir, @@ -185,6 +229,8 @@ func (r *Runner) runMassdns(inputFile string) { StrictWildcard: r.options.StrictWildcard, WildcardOutputFile: r.options.WildcardOutputFile, MassDnsCmd: r.options.MassDnsCmd, + KeepStderr: r.options.KeepStderr, + BatchSize: r.options.BatchSize, OnResult: r.options.OnResult, }) if err != nil { @@ -194,12 +240,13 @@ func (r *Runner) runMassdns(inputFile string) { err = massdns.Run(context.Background()) if err != nil { - gologger.Error().Msgf("Could not run massdns: %s\n", err) + gologger.Error().Msgf("Could not process existing massdns output: %s\n", err) + return } if r.options.WildcardOutputFile != "" { _ = massdns.DumpWildcardsToFile(r.options.WildcardOutputFile) } - gologger.Info().Msgf("Finished resolving.\n") + gologger.Info().Msgf("Finished processing existing output.\n") } diff --git a/pkg/store/store.go b/pkg/store/store.go index 2632b44..c6c1aed 100644 --- a/pkg/store/store.go +++ b/pkg/store/store.go @@ -1,10 +1,11 @@ package store import ( + "encoding/json" "os" "strings" - sliceutil "github.com/projectdiscovery/utils/slice" + mapsutil "github.com/projectdiscovery/utils/maps" "github.com/syndtr/goleveldb/leveldb" "github.com/syndtr/goleveldb/leveldb/opt" ) @@ -23,7 +24,13 @@ func New(dbPath string) (*Store, error) { return nil, err } db, err := leveldb.OpenFile(storeDb, &opt.Options{ - CompactionTableSize: 256 * Megabyte, + // Optimize for disk space reduction + CompactionTableSize: 64 * Megabyte, // Reduced from 256MB for more frequent compaction + WriteBuffer: 2 * Megabyte, // Reduced from 4MB for more frequent flushing + WriteL0SlowdownTrigger: 4, // Trigger slowdown earlier + WriteL0PauseTrigger: 8, // Trigger pause earlier + BlockSize: 2 * 1024, // Reduced from 4KB for better compression of small records + BlockCacheCapacity: 4 * Megabyte, // Reduced from 8MB to lower memory usage }) if err != nil { return nil, err @@ -33,7 +40,12 @@ func New(dbPath string) (*Store, error) { // New creates a new ip-hostname pair in the map func (s *Store) New(ip, hostname string) error { - return s.DB.Put([]byte(ip), []byte(hostname), nil) + hostnameMap := map[string]struct{}{hostname: {}} + jsonData, err := json.Marshal(hostnameMap) + if err != nil { + return err + } + return s.DB.Put([]byte(ip), jsonData, nil) } // Exists indicates if an IP exists in the map @@ -44,20 +56,44 @@ func (s *Store) Exists(ip string) bool { // Get gets the meta-information for an IP address from the map. func (s *Store) GetHostnames(ip string) string { - hostname, err := s.DB.Get([]byte(ip), nil) + data, err := s.DB.Get([]byte(ip), nil) if err != nil { return "" } - return string(hostname) + + var hostnameMap map[string]struct{} + if err := json.Unmarshal(data, &hostnameMap); err != nil { + return "" + } + + return strings.Join(mapsutil.GetKeys(hostnameMap), ",") } func (s *Store) Append(ip string, hostnames ...string) error { - existingHostnames, _ := s.DB.Get([]byte(ip), nil) - if len(existingHostnames) > 0 { - hostnames = append(hostnames, string(existingHostnames)) + // Get existing hostnames + var hostnameMap map[string]struct{} + existingData, err := s.DB.Get([]byte(ip), nil) + if err == nil && len(existingData) > 0 { + if err := json.Unmarshal(existingData, &hostnameMap); err != nil { + // If unmarshaling fails, start with empty map + hostnameMap = make(map[string]struct{}) + } + } else { + hostnameMap = make(map[string]struct{}) + } + + // Add new hostnames to map (automatic deduplication) + for _, hostname := range hostnames { + hostnameMap[hostname] = struct{}{} + } + + // Marshal and store + jsonData, err := json.Marshal(hostnameMap) + if err != nil { + return err } - return s.DB.Put([]byte(ip), []byte(strings.Join(hostnames, ",")), nil) + return s.DB.Put([]byte(ip), jsonData, nil) } // Delete deletes the records for an IP from store. @@ -75,8 +111,18 @@ func (s *Store) Iterate(f func(ip string, hostnames []string, counter int)) { for iter.Next() { ip := string(iter.Key()) - hostnames := strings.Split(string(iter.Value()), ",") - hostnames = sliceutil.Dedupe(hostnames) + + var hostnameMap map[string]struct{} + if err := json.Unmarshal(iter.Value(), &hostnameMap); err != nil { + continue + } + + // Convert map keys to slice + hostnames := make([]string, 0, len(hostnameMap)) + for hostname := range hostnameMap { + hostnames = append(hostnames, hostname) + } + counter := len(hostnames) f(ip, hostnames, counter) }