main.go

package main

import (
	"bufio"
	"crypto/tls"
	"flag"
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"

	"strings"

	"github.com/fatih/color"
)

// Array flags for multiple headers
type headerFlags []string

func (h *headerFlags) String() string {
	return fmt.Sprint(*h)
}

func (h *headerFlags) Set(value string) error {
	*h = append(*h, value)
	return nil
}

func main() {
	// CLI Flags
	var headers headerFlags
	concurrency := flag.Int("w", 20, "Number of concurrent workers")
	singleTarget := flag.String("u", "", "Single target URL/Domain to scan")
	verbose := flag.Bool("v", false, "Enable verbose output")
	flag.Var(&headers, "H", "Custom header (can be used multiple times, e.g. -H 'Authorization: Token')")

	// Custom Usage
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: cat domains.txt | ghostpack [options]\n")
		fmt.Fprintf(os.Stderr, "       ghostpack -u https://target.com [options]\n\n")
		fmt.Fprintf(os.Stderr, "GhostPack - Dependency Confusion Scanner\n\n")
		fmt.Fprintf(os.Stderr, "Options:\n")
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\nExample:\n")
		fmt.Fprintf(os.Stderr, "  cat domains.txt | ghostpack -w 50 -H \"Cookie: session=123\"\n")
	}

	flag.Parse()

	// Input reading
	var domains []string

	if *singleTarget != "" {
		domains = append(domains, *singleTarget)
	} else {
		stat, _ := os.Stdin.Stat()
		if (stat.Mode() & os.ModeCharDevice) != 0 {
			flag.Usage()
			os.Exit(0)
		}

		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			if line != "" {
				domains = append(domains, line)
			}
		}
	}

	if len(domains) == 0 {
		fmt.Println("No domains provided.")
		os.Exit(1)
	}

	// Output styling
	red := color.New(color.FgRed, color.Bold).SprintFunc()
	green := color.New(color.FgGreen).SprintFunc()
	yellow := color.New(color.FgYellow).SprintFunc()
	cyan := color.New(color.FgCyan).SprintFunc()

	fmt.Printf("GhostPack started. Scanning %d domains with %d workers...\n", len(domains), *concurrency)

	// HTTP Clients setup
	baseTransport := &http.Transport{
		DisableKeepAlives: true,
		TLSClientConfig:   &tls.Config{InsecureSkipVerify: true},
	}

	// Analysis Client (Clean)
	analysisClient := &http.Client{
		Timeout:   10 * time.Second,
		Transport: baseTransport,
	}

	// Scan Client (With Headers)
	// Note: We need a fresh transport for headers that wraps the base transport
	// But `TransportWithHeaders` wraps `http.RoundTripper`.
	scanTransport := &TransportWithHeaders{
		Wrapped: baseTransport,
		Headers: headers,
	}

	scanClient := &http.Client{
		Timeout:   10 * time.Second,
		Transport: scanTransport,
	}

	// Logger
	var log LogFunc
	if *verbose {
		log = func(format string, a ...interface{}) {
			fmt.Printf(format+"\n", a...)
		}
	}

	// Worker Pool Implementation
	jobs := make(chan string, len(domains))
	results := make(chan string, len(domains)) // Buffer to avoid blocking workers if writer is slow
	var wg sync.WaitGroup

	// Start Workers
	for i := 0; i < *concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for domain := range jobs {
				output := processDomain(domain, scanClient, analysisClient, log, red, green, yellow, cyan)
				if output != "" {
					results <- output
				}
			}
		}()
	}

	// Start Writer (Single goroutine for thread-safe output)
	done := make(chan struct{})
	go func() {
		for msg := range results {
			fmt.Print(msg)
		}
		close(done)
	}()

	// Push jobs directly (since we have all domains in memory)
	for _, d := range domains {
		jobs <- d
	}
	close(jobs)

	// Wait for workers
	wg.Wait()
	close(results)

	// Wait for writer
	<-done
}

// TransportWithHeaders adds custom headers
type TransportWithHeaders struct {
	Wrapped http.RoundTripper
	Headers []string
}

func (t *TransportWithHeaders) RoundTrip(req *http.Request) (*http.Response, error) {
	r2 := new(http.Request)
	*r2 = *req
	r2.Header = make(http.Header, len(req.Header))
	for k, s := range req.Header {
		r2.Header[k] = append([]string(nil), s...)
	}

	for _, h := range t.Headers {
		parts := strings.SplitN(h, ":", 2)
		if len(parts) == 2 {
			r2.Header.Set(strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1]))
		}
	}

	return t.Wrapped.RoundTrip(r2)
}

// processDomain now returns a string instead of printing directly
func processDomain(domain string, scanClient, analysisClient *http.Client, log LogFunc, red, green, yellow, cyan func(a ...interface{}) string) string {
	var sb strings.Builder

	files, err := ProbeDomain(domain, scanClient, log)
	if err != nil {
		return ""
	}

	analyzedPackages := make(map[string]bool)

	for _, file := range files {
		var packages []string
		var parseErr error

		if file.Type == "json" {
			packages, parseErr = ParsePackageJSON(file.Content)
		} else if file.Type == "yarn" {
			packages, parseErr = ParseYarnLock(file.Content)
		}

		if parseErr != nil {
			continue
		}

		sb.WriteString(fmt.Sprintf("%s %s | Found %s\n", green("[FOUND]"), domain, file.URL))

		for _, pkg := range packages {
			if active, exists := analyzedPackages[pkg]; exists && active {
				continue
			}
			analyzedPackages[pkg] = true

			result := CheckPackage(analysisClient, pkg)

			if result.Status == "VULNERABLE" {
				prefix := ""
				if result.IsScoped {
					prefix = red("[!!! SCOPED !!!] ")
				}
				sb.WriteString(fmt.Sprintf("%s%s %s | %s | Package: %s\n",
					prefix,
					red("[VULNERABLE]"),
					domain,
					file.URL,
					red(pkg),
				))
			} else if result.Status == "SAFE" {
				if result.IsScoped {
					sb.WriteString(fmt.Sprintf("%s %s | %s | Package: %s\n",
						green("[SAFE]"),
						domain,
						file.URL,
						pkg,
					))
				}
			}
		}
	}
	return sb.String()
}