god-eye/internal/nucleitpl/download.go

package nucleitpl

import (
	"archive/zip"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"time"
)

// TemplatesZipURL is the default ZIP archive of the projectdiscovery
// nuclei-templates repository (main branch).
const TemplatesZipURL = "https://github.com/projectdiscovery/nuclei-templates/archive/refs/heads/main.zip"

// Downloader fetches the nuclei-templates archive and extracts the
// YAML files into destDir. Designed to be invoked at most once per
// scan: after a successful extraction the destination dir persists
// across runs; subsequent invocations return quickly via hasTemplates().
type Downloader struct {
	// ZipURL overrides TemplatesZipURL for testing or mirroring.
	ZipURL string
	// HTTPClient is used for the download. Default: 10-minute timeout.
	HTTPClient *http.Client
	// Writer receives progress lines when Verbose is true. Defaults to
	// os.Stderr.
	Writer io.Writer
	// Verbose toggles progress logging.
	Verbose bool
	// MinTemplatesToConsiderPresent is the count of .yaml files under
	// destDir below which we treat the directory as empty / incomplete
	// and re-download. Default: 50.
	MinTemplatesToConsiderPresent int
}

// NewDownloader returns a Downloader with sensible defaults.
func NewDownloader() *Downloader {
	return &Downloader{
		ZipURL:                        TemplatesZipURL,
		HTTPClient:                    &http.Client{Timeout: 10 * time.Minute},
		Writer:                        os.Stderr,
		MinTemplatesToConsiderPresent: 50,
	}
}

// EnsureTemplates guarantees destDir contains a usable set of Nuclei
// YAML templates. If the directory already has ≥ MinTemplatesToConsiderPresent
// templates, it's a no-op. Otherwise the ZIP is downloaded, streamed to
// a temp file, and extracted (YAML files only).
//
// destDir is created if it doesn't exist.
func (d *Downloader) EnsureTemplates(destDir string) error {
	if destDir == "" {
		return errors.New("EnsureTemplates: empty destDir")
	}
	if d.hasEnoughTemplates(destDir) {
		if d.Verbose {
			fmt.Fprintf(d.writer(), "✓ nuclei templates already present at %s\n", destDir)
		}
		return nil
	}

	if err := os.MkdirAll(destDir, 0o755); err != nil {
		return fmt.Errorf("mkdir %s: %w", destDir, err)
	}

	if d.Verbose {
		fmt.Fprintf(d.writer(), "↓ downloading nuclei-templates from %s\n", d.zipURL())
	}

	tmpPath, err := d.downloadZip()
	if err != nil {
		return err
	}
	defer os.Remove(tmpPath)

	count, bytes, err := d.extractYAML(tmpPath, destDir)
	if err != nil {
		return err
	}

	if count < d.MinTemplatesToConsiderPresent {
		return fmt.Errorf("extracted only %d templates (expected ≥ %d) — archive may be incomplete", count, d.MinTemplatesToConsiderPresent)
	}

	if d.Verbose {
		fmt.Fprintf(d.writer(), "✓ extracted %d nuclei templates (%s) into %s\n",
			count, humanBytesN(bytes), destDir)
	}
	return nil
}

// Refresh forces a re-download regardless of current directory contents.
// Useful for `god-eye nuclei-update` style CLI commands.
func (d *Downloader) Refresh(destDir string) error {
	if destDir == "" {
		return errors.New("Refresh: empty destDir")
	}
	if err := os.MkdirAll(destDir, 0o755); err != nil {
		return fmt.Errorf("mkdir %s: %w", destDir, err)
	}

	if d.Verbose {
		fmt.Fprintf(d.writer(), "↓ refreshing nuclei-templates from %s\n", d.zipURL())
	}

	tmpPath, err := d.downloadZip()
	if err != nil {
		return err
	}
	defer os.Remove(tmpPath)

	count, bytes, err := d.extractYAML(tmpPath, destDir)
	if err != nil {
		return err
	}
	if d.Verbose {
		fmt.Fprintf(d.writer(), "✓ refreshed %d templates (%s)\n", count, humanBytesN(bytes))
	}
	return nil
}

// --- internals -----------------------------------------------------------

func (d *Downloader) hasEnoughTemplates(dir string) bool {
	info, err := os.Stat(dir)
	if err != nil || !info.IsDir() {
		return false
	}
	found := 0
	threshold := d.MinTemplatesToConsiderPresent
	if threshold <= 0 {
		threshold = 50
	}
	_ = filepath.Walk(dir, func(_ string, fi os.FileInfo, err error) error {
		if err != nil {
			return nil
		}
		if fi.IsDir() {
			return nil
		}
		name := strings.ToLower(fi.Name())
		if strings.HasSuffix(name, ".yaml") || strings.HasSuffix(name, ".yml") {
			found++
			if found >= threshold {
				return filepath.SkipAll
			}
		}
		return nil
	})
	return found >= threshold
}

func (d *Downloader) zipURL() string {
	if d.ZipURL != "" {
		return d.ZipURL
	}
	return TemplatesZipURL
}

func (d *Downloader) writer() io.Writer {
	if d.Writer != nil {
		return d.Writer
	}
	return os.Stderr
}

func (d *Downloader) downloadZip() (string, error) {
	client := d.HTTPClient
	if client == nil {
		client = &http.Client{Timeout: 10 * time.Minute}
	}

	req, err := http.NewRequest("GET", d.zipURL(), nil)
	if err != nil {
		return "", err
	}
	req.Header.Set("User-Agent", "god-eye-v2")
	req.Header.Set("Accept", "application/zip")

	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("download: %w", err)
	}
	defer resp.Body.Close()

	// Follow standard HTTP error reporting.
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("download: HTTP %d from %s", resp.StatusCode, d.zipURL())
	}

	tmp, err := os.CreateTemp("", "nuclei-templates-*.zip")
	if err != nil {
		return "", fmt.Errorf("create temp: %w", err)
	}

	// Streaming copy with throttled progress output.
	var written atomic.Int64
	pr := &progressReader{
		r:       resp.Body,
		written: &written,
		verbose: d.Verbose,
		writer:  d.writer(),
		total:   resp.ContentLength,
		prefix:  "  downloading",
	}

	if _, err := io.Copy(tmp, pr); err != nil {
		tmp.Close()
		os.Remove(tmp.Name())
		return "", fmt.Errorf("stream download: %w", err)
	}
	if err := tmp.Close(); err != nil {
		os.Remove(tmp.Name())
		return "", err
	}
	return tmp.Name(), nil
}

// extractYAML walks the zip and writes every .yaml / .yml file into
// destDir. Returns (count, totalBytes, error).
//
// The top-level directory in the archive (e.g. "nuclei-templates-main/")
// is stripped so entries land at destDir/<category>/<file>.yaml.
//
// Path-traversal protection: every resolved destination must be within
// destDir; otherwise the entry is skipped.
func (d *Downloader) extractYAML(zipPath, destDir string) (int, int64, error) {
	zr, err := zip.OpenReader(zipPath)
	if err != nil {
		return 0, 0, fmt.Errorf("open zip: %w", err)
	}
	defer zr.Close()

	absDest, err := filepath.Abs(destDir)
	if err != nil {
		return 0, 0, err
	}

	var count int
	var bytes int64
	for _, f := range zr.File {
		if f.FileInfo().IsDir() {
			continue
		}
		lower := strings.ToLower(f.Name)
		if !strings.HasSuffix(lower, ".yaml") && !strings.HasSuffix(lower, ".yml") {
			continue
		}

		// Strip leading top-level folder if present.
		rel := f.Name
		if i := strings.Index(rel, "/"); i >= 0 {
			rel = rel[i+1:]
		}
		if rel == "" {
			continue
		}
		// Guard against path traversal / absolute paths.
		if strings.Contains(rel, "..") || filepath.IsAbs(rel) {
			continue
		}

		dest := filepath.Join(absDest, rel)
		if !strings.HasPrefix(dest, absDest+string(os.PathSeparator)) && dest != absDest {
			continue
		}

		if err := os.MkdirAll(filepath.Dir(dest), 0o755); err != nil {
			continue
		}

		rc, err := f.Open()
		if err != nil {
			continue
		}
		out, err := os.OpenFile(dest, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644)
		if err != nil {
			rc.Close()
			continue
		}
		n, cerr := io.Copy(out, rc)
		rc.Close()
		out.Close()
		if cerr != nil {
			_ = os.Remove(dest)
			continue
		}
		count++
		bytes += n
	}
	return count, bytes, nil
}

// --- helpers -------------------------------------------------------------

// progressReader wraps an io.Reader and emits throttled progress lines
// as bytes are consumed. Throttling: one line every ~5% of total (or
// every ~5MB when total is unknown).
type progressReader struct {
	r       io.Reader
	written *atomic.Int64
	total   int64
	verbose bool
	writer  io.Writer
	prefix  string

	lastPct    int
	lastBytes  int64
	lastReport time.Time
}

func (p *progressReader) Read(b []byte) (int, error) {
	n, err := p.r.Read(b)
	if n > 0 {
		p.written.Add(int64(n))
		if p.verbose {
			p.maybeReport()
		}
	}
	return n, err
}

func (p *progressReader) maybeReport() {
	w := p.written.Load()

	// Rate-limit prints to avoid flooding the terminal.
	if time.Since(p.lastReport) < 200*time.Millisecond {
		return
	}

	if p.total > 0 {
		pct := int(float64(w) / float64(p.total) * 100)
		if pct >= p.lastPct+5 || pct == 100 {
			fmt.Fprintf(p.writer, "%s %3d%%  %s / %s\n",
				p.prefix, pct, humanBytesN(w), humanBytesN(p.total))
			p.lastPct = pct
			p.lastReport = time.Now()
		}
	} else {
		// Unknown total: report every ~5MB.
		if w-p.lastBytes >= 5*1024*1024 {
			fmt.Fprintf(p.writer, "%s %s\n", p.prefix, humanBytesN(w))
			p.lastBytes = w
			p.lastReport = time.Now()
		}
	}
}

// humanBytesN formats a byte count like "2.3MB". Duplicated from
// ai/ensure.go to avoid a cross-package dependency.
func humanBytesN(n int64) string {
	const k = 1024.0
	if n < int64(k) {
		return fmt.Sprintf("%dB", n)
	}
	units := []string{"KB", "MB", "GB", "TB"}
	v := float64(n) / k
	for _, u := range units {
		if v < k {
			return fmt.Sprintf("%.1f%s", v, u)
		}
		v /= k
	}
	return fmt.Sprintf("%.1fPB", v)
}