main.go

package main

import (
	"bufio"
	"context"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/NullifiedSec/voidprobber/pkg/config"
	"github.com/NullifiedSec/voidprobber/pkg/output"
	"github.com/NullifiedSec/voidprobber/pkg/probe"
	"github.com/NullifiedSec/voidprobber/pkg/screenshot"
	"github.com/NullifiedSec/voidprobber/pkg/worker"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var (
	version = "2.0.0"
	cfg     *config.Config
	logger  *logrus.Logger
)

func main() {
	if err := rootCmd.Execute(); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}
}

var rootCmd = &cobra.Command{
	Use:   "voidprobber",
	Short: "Take a list of domains and probe for working HTTP and HTTPS servers",
	Long: `voidprobber takes a list of domains and probes for working HTTP and HTTPS servers.
It supports various output formats, rate limiting, retries, and advanced configuration options.`,
	Version: version,
	RunE:    run,
}

func init() {
	cobra.OnInitialize(initConfig)

	// Global flags
	rootCmd.PersistentFlags().String("config", "", "config file (default is ./voidprobber.yaml)")
	rootCmd.PersistentFlags().BoolP("verbose", "v", false, "verbose output")
	rootCmd.PersistentFlags().BoolP("quiet", "q", false, "quiet mode (no progress or stats)")
	
	// Core flags
	rootCmd.Flags().IntP("concurrency", "c", 20, "concurrency level")
	rootCmd.Flags().DurationP("timeout", "t", 10*time.Second, "request timeout")
	rootCmd.Flags().Duration("connect-timeout", 5*time.Second, "connection timeout")
	rootCmd.Flags().Duration("read-timeout", 10*time.Second, "read timeout")
	
	// Probe flags
	rootCmd.Flags().StringSliceP("probe", "p", []string{}, "additional probe (proto:port or predefined list)")
	rootCmd.Flags().IntSlice("ports", []int{}, "additional ports to probe")
	rootCmd.Flags().Bool("ports-common", false, "scan common ports (top 1000)")
	rootCmd.Flags().BoolP("skip-default", "s", false, "skip default probes (http:80 and https:443)")
	rootCmd.Flags().Bool("prefer-https", false, "only try HTTP if HTTPS fails")
	
	// HTTP flags
	rootCmd.Flags().StringP("method", "m", "GET", "HTTP method")
	rootCmd.Flags().String("user-agent", "voidprobber/2.0", "User-Agent header")
	rootCmd.Flags().StringSlice("header", []string{}, "custom headers (key:value)")
	rootCmd.Flags().String("proxy", "", "proxy URL (http://proxy:8080 or socks5://proxy:1080)")
	rootCmd.Flags().Bool("follow-redirects", false, "follow redirects")
	rootCmd.Flags().Int("max-redirects", 5, "maximum redirects to follow")
	
	// TLS flags
	rootCmd.Flags().Bool("insecure", true, "skip TLS certificate verification")
	rootCmd.Flags().String("tls-server-name", "", "TLS server name for verification")
	rootCmd.Flags().Bool("cert-info", false, "include certificate information in output")
	
	// Output flags
	rootCmd.Flags().StringP("output", "o", "", "output file")
	rootCmd.Flags().String("format", "text", "output format (text, json, csv)")
	rootCmd.Flags().Bool("show-progress", false, "show progress")
	rootCmd.Flags().Bool("show-stats", true, "show final statistics")
	rootCmd.Flags().Bool("show-status-codes", false, "show HTTP status codes in text output")
	rootCmd.Flags().Bool("color", true, "colorize output (use --no-color to disable)")
	rootCmd.Flags().Bool("no-color", false, "disable colored output")
	
	// Rate limiting flags
	rootCmd.Flags().Int("rate-limit", 0, "requests per second (0 = unlimited)")
	rootCmd.Flags().Duration("rate-limit-per", time.Second, "rate limit time window")
	
	// Retry flags
	rootCmd.Flags().Int("max-retries", 0, "maximum retries per probe")
	rootCmd.Flags().Duration("retry-delay", time.Second, "delay between retries")
	rootCmd.Flags().Float64("retry-backoff", 1.5, "retry backoff multiplier")
	
	// Advanced flags
	rootCmd.Flags().Bool("http2-only", false, "only use HTTP/2")
	rootCmd.Flags().Bool("check-content", false, "analyze response content (extract title)")
	rootCmd.Flags().Bool("title", false, "extract and display page titles")
	rootCmd.Flags().IntSlice("status-codes", []int{}, "consider only these status codes as success")
	
	// Screenshot flags
	rootCmd.Flags().Bool("screenshot", false, "capture screenshots of discovered services")
	rootCmd.Flags().String("screenshot-dir", "screenshots", "directory to save screenshots")
	rootCmd.Flags().Int("screenshot-width", 1280, "screenshot width in pixels")
	rootCmd.Flags().Int("screenshot-height", 720, "screenshot height in pixels")
	rootCmd.Flags().Duration("screenshot-timeout", 30*time.Second, "screenshot timeout")
	rootCmd.Flags().Bool("screenshot-fullpage", false, "capture full page screenshot")
	
	// Bind flags to viper
	viper.BindPFlags(rootCmd.Flags())
	viper.BindPFlags(rootCmd.PersistentFlags())
}

func initConfig() {
	// Initialize logger
	logger = logrus.New()
	logger.SetFormatter(&logrus.TextFormatter{
		FullTimestamp: true,
	})

	// Set log level based on flags
	if viper.GetBool("verbose") {
		logger.SetLevel(logrus.DebugLevel)
	} else if viper.GetBool("quiet") {
		logger.SetLevel(logrus.ErrorLevel)
	} else {
		logger.SetLevel(logrus.InfoLevel)
	}

	// Load configuration
	configFile := viper.GetString("config")
	var err error
	cfg, err = config.LoadConfig(configFile)
	if err != nil {
		logger.Fatalf("Failed to load config: %v", err)
	}

	// Override config with command line flags
	overrideConfigFromFlags()
}

func overrideConfigFromFlags() {
	if viper.IsSet("concurrency") {
		cfg.Concurrency = viper.GetInt("concurrency")
	}
	if viper.IsSet("timeout") {
		cfg.Timeout = viper.GetDuration("timeout")
	}
	if viper.IsSet("connect-timeout") {
		cfg.ConnectTimeout = viper.GetDuration("connect-timeout")
	}
	if viper.IsSet("read-timeout") {
		cfg.ReadTimeout = viper.GetDuration("read-timeout")
	}
	if viper.IsSet("probe") {
		cfg.Probes = viper.GetStringSlice("probe")
	}
	if viper.IsSet("ports") {
		cfg.CustomPorts = viper.GetIntSlice("ports")
	}
	if viper.IsSet("ports-common") {
		cfg.ScanCommonPorts = viper.GetBool("ports-common")
	}
	if viper.IsSet("skip-default") {
		cfg.SkipDefault = viper.GetBool("skip-default")
	}
	if viper.IsSet("prefer-https") {
		cfg.PreferHTTPS = viper.GetBool("prefer-https")
	}
	if viper.IsSet("method") {
		cfg.Method = viper.GetString("method")
	}
	if viper.IsSet("user-agent") {
		cfg.UserAgent = viper.GetString("user-agent")
	}
	if viper.IsSet("header") {
		headers := viper.GetStringSlice("header")
		cfg.Headers = make(map[string]string)
		for _, header := range headers {
			parts := strings.SplitN(header, ":", 2)
			if len(parts) == 2 {
				cfg.Headers[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
			}
		}
	}
	if viper.IsSet("proxy") {
		cfg.ProxyURL = viper.GetString("proxy")
	}
	if viper.IsSet("follow-redirects") {
		cfg.FollowRedirects = viper.GetBool("follow-redirects")
	}
	if viper.IsSet("max-redirects") {
		cfg.MaxRedirects = viper.GetInt("max-redirects")
	}
	if viper.IsSet("insecure") {
		cfg.InsecureTLS = viper.GetBool("insecure")
	}
	if viper.IsSet("tls-server-name") {
		cfg.TLSServerName = viper.GetString("tls-server-name")
	}
	if viper.IsSet("cert-info") {
		cfg.CertInfo = viper.GetBool("cert-info")
	}
	if viper.IsSet("output") {
		cfg.OutputFile = viper.GetString("output")
	}
	if viper.IsSet("format") {
		cfg.OutputFormat = viper.GetString("format")
	}
	if viper.IsSet("show-progress") {
		cfg.ShowProgress = viper.GetBool("show-progress")
	}
	if viper.IsSet("show-stats") {
		cfg.ShowStats = viper.GetBool("show-stats")
	}
	if viper.IsSet("show-status-codes") {
		cfg.ShowStatusCodes = viper.GetBool("show-status-codes")
	}
	
	// Handle color flags (--no-color overrides --color)
	if viper.IsSet("no-color") && viper.GetBool("no-color") {
		cfg.ColorOutput = false
	} else if viper.IsSet("color") {
		cfg.ColorOutput = viper.GetBool("color")
	}
	if viper.IsSet("rate-limit") {
		cfg.RateLimit = viper.GetInt("rate-limit")
	}
	if viper.IsSet("rate-limit-per") {
		cfg.RateLimitPer = viper.GetDuration("rate-limit-per")
	}
	if viper.IsSet("max-retries") {
		cfg.MaxRetries = viper.GetInt("max-retries")
	}
	if viper.IsSet("retry-delay") {
		cfg.RetryDelay = viper.GetDuration("retry-delay")
	}
	if viper.IsSet("retry-backoff") {
		cfg.RetryBackoff = viper.GetFloat64("retry-backoff")
	}
	if viper.IsSet("http2-only") {
		cfg.HTTP2Only = viper.GetBool("http2-only")
	}
	if viper.IsSet("check-content") {
		cfg.CheckContent = viper.GetBool("check-content")
	}
	if viper.IsSet("title") {
		cfg.ExtractTitle = viper.GetBool("title")
	}
	if viper.IsSet("status-codes") {
		cfg.StatusCodes = viper.GetIntSlice("status-codes")
	}
	
	// Screenshot configuration
	if viper.IsSet("screenshot") {
		cfg.TakeScreenshots = viper.GetBool("screenshot")
	}
	if viper.IsSet("screenshot-dir") {
		cfg.ScreenshotDir = viper.GetString("screenshot-dir")
	}
	if viper.IsSet("screenshot-width") {
		cfg.ScreenshotWidth = viper.GetInt("screenshot-width")
	}
	if viper.IsSet("screenshot-height") {
		cfg.ScreenshotHeight = viper.GetInt("screenshot-height")
	}
	if viper.IsSet("screenshot-timeout") {
		cfg.ScreenshotTimeout = viper.GetDuration("screenshot-timeout")
	}
	if viper.IsSet("screenshot-fullpage") {
		cfg.ScreenshotFullPage = viper.GetBool("screenshot-fullpage")
	}
	
	if viper.IsSet("verbose") {
		cfg.Verbose = viper.GetBool("verbose")
	}
	if viper.IsSet("quiet") {
		cfg.Quiet = viper.GetBool("quiet")
	}
}

func run(cmd *cobra.Command, args []string) error {
	// Create context for cancellation
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Handle signals
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	
	// Create screenshot service
	screenshotService, err := screenshot.NewService(cfg, logger)
	if err != nil {
		logger.WithError(err).Warn("Failed to initialize screenshot service, continuing without screenshots")
		screenshotService = nil
	}
	if screenshotService != nil {
		defer screenshotService.Close()
	}

	// Create prober
	prober := probe.New(cfg, logger, screenshotService)

	// Create worker pool
	pool := worker.NewPool(cfg, prober, logger)

	// Create output writer
	writer, err := output.NewWriter(cfg.OutputFormat, cfg.OutputFile, cfg.ExtractTitle, cfg.ShowStatusCodes, cfg.ColorOutput, logger)
	if err != nil {
		return fmt.Errorf("failed to create output writer: %w", err)
	}
	defer writer.Close()

	// Create verbose output if needed
	var verboseOutput *output.VerboseOutput
	if cfg.Verbose {
		verboseOutput = output.NewVerboseOutput(logger)
	}

	// Create stats reporter
	statsReporter := output.NewStatsReporter(logger, cfg.Quiet)

	// Start worker pool
	pool.Start()

	// Read domains from stdin
	domains := make(chan string, 100)
		go func() {
		defer close(domains)
		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			domain := strings.TrimSpace(scanner.Text())
			if domain != "" {
				select {
				case domains <- domain:
				case <-ctx.Done():
					return
				}
			}
		}
		if err := scanner.Err(); err != nil {
			logger.WithError(err).Error("Error reading from stdin")
		}
	}()

	// Generate and submit jobs
	var jobID int64
	var totalDomains int64
	jobsDone := make(chan struct{})

	go func() {
		defer close(jobsDone)
		for domain := range domains {
			select {
			case <-ctx.Done():
				return
			default:
			}
			
			totalDomains++
			targets := probe.GenerateTargets(domain, cfg)
			
			for _, target := range targets {
				select {
				case <-ctx.Done():
					return
				default:
				}
				
				jobID++
				job := worker.Job{
					URL:    target,
					Domain: domain,
					ID:     jobID,
				}
				pool.AddJob(job)
			}
		}
	}()

	// Setup progress tracking
	var progressTracker *worker.ProgressTracker
	if cfg.ShowProgress {
		progressTracker = worker.NewProgressTracker(0, logger)
		progressTracker.Start(5 * time.Second)
		defer progressTracker.Stop()
	}

	// Process results
	var processedResults int64
	resultsDone := make(chan struct{})

	go func() {
		defer close(resultsDone)
		
		for {
			select {
			case result, ok := <-pool.Results():
				if !ok {
					// Channel closed, we're done
					return
				}
				
				processedResults++
				
				// Write verbose output
				if verboseOutput != nil {
					verboseOutput.WriteVerboseResult(result)
				}
				
				// Write result
				if err := writer.WriteResult(result); err != nil {
					logger.WithError(err).Error("Failed to write result")
				}
				
				// Update progress
				if progressTracker != nil {
					progressTracker.Increment()
				}
				
			case <-ctx.Done():
				// Context cancelled, stop processing
				return
			}
		}
	}()

	// Main coordination loop
	done := make(chan struct{})
	go func() {
		defer close(done)
		
		// Wait for jobs to be submitted
		select {
		case <-jobsDone:
			logger.Debug("All jobs submitted")
		case <-ctx.Done():
			logger.Debug("Context cancelled while waiting for jobs")
			return
		}
		
		// Stop the worker pool (this will close the results channel when all work is done)
		pool.Stop()
		
		// Wait for all results to be processed
		select {
		case <-resultsDone:
			logger.Debug("All results processed")
		case <-ctx.Done():
			logger.Debug("Context cancelled while waiting for results")
		}
	}()

	// Wait for completion or cancellation
	select {
	case <-done:
		// Normal completion
		logger.Debug("Scan completed normally")
	case <-sigChan:
		logger.Info("Received interrupt signal, shutting down...")
		cancel()
		
		// Give a brief moment for graceful shutdown
		select {
		case <-done:
			logger.Debug("Graceful shutdown completed")
		case <-time.After(2 * time.Second):
			logger.Warn("Forced shutdown after timeout")
		}
	case <-ctx.Done():
		logger.Debug("Context cancelled")
	}

	// Ensure everything is cleaned up
	cancel()
	
	// Stop pool if not already stopped
	pool.Stop()

	// Report final statistics
	if cfg.ShowStats {
		stats := pool.GetStats()
		statsReporter.ReportFinalStats(&stats, totalDomains)
	}

	logger.Info("Scan completed")
	return nil
}