// This file is part of checker-dnsviz.
//
// checker-dnsviz is free software: you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the Free
// Software Foundation, version 2.
//
// checker-dnsviz is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
// or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
// more details.
//
// You should have received a copy of the GNU General Public License along with
// checker-dnsviz. If not, see <https://www.gnu.org/licenses/>.
//
// SPDX-License-Identifier: GPL-2.0-only

// Package collect contains the DNSViz subprocess invocation. It is kept
// separate from the checker package so that the checker package (pure analysis
// logic) can be imported under MIT terms without pulling in GPL-covered code.
package collect

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"
	"unicode/utf8"

	checker "git.happydns.org/checker-dnsviz/checker"
	sdk "git.happydns.org/checker-sdk-go/checker"
)

// defaultTrustAnchorsFile is the BIND-formatted root DNSKEY trust anchor file
// shipped by Alpine's `dnssec-root` package (installed in our Docker image).
// Passed to `dnsviz grok -t` so the root zone gets classified as SECURE
// instead of staying at the DNS-rcode "NOERROR" fallback.
const defaultTrustAnchorsFile = "/usr/share/dnssec-root/trusted-key.key"

const (
	defaultProbeTimeout  = 120 * time.Second
	maxDNSVizOutputBytes = 16 << 20 // 16 MiB
)

// Collector holds the runtime configuration for DNSViz invocations.
type Collector struct {
	// Bin is the path to the dnsviz CLI. Defaults to "dnsviz".
	Bin string
	// ExtraArgs is a whitespace-separated list of extra arguments appended to
	// `dnsviz probe`. Defaults to "-A".
	ExtraArgs string
	// TrustAnchorsFile is a BIND-format DNSKEY file used as DNSSEC trust
	// anchor (passed to `dnsviz grok -t`). When empty, defaultTrustAnchorsFile
	// is used if it exists; otherwise grok runs without `-t` and the root
	// zone falls back to a NOERROR classification.
	TrustAnchorsFile string
}

// Collect runs `dnsviz probe | dnsviz grok` against the domain named in opts
// and returns the structured analysis as a *checker.DNSVizData.
func (c *Collector) Collect(ctx context.Context, opts sdk.CheckerOptions) (any, error) {
	domain, _ := sdk.GetOption[string](opts, "domain_name")
	domain = strings.TrimSpace(strings.TrimSuffix(domain, "."))
	if domain == "" {
		return nil, fmt.Errorf("missing 'domain_name' option")
	}
	if !isValidDomainName(domain) {
		return nil, fmt.Errorf("invalid 'domain_name' option")
	}

	timeout := defaultProbeTimeout
	if n := sdk.GetIntOption(opts, "probeTimeoutSeconds", 0); n > 0 {
		timeout = time.Duration(n) * time.Second
	}

	bin := strings.TrimSpace(c.Bin)
	if bin == "" {
		bin = "dnsviz"
	}
	extraArgs := c.ExtraArgs
	if extraArgs == "" {
		extraArgs = "-A"
	}

	probeCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	// Probe the queried name AND every ancestor up to the root. dnsviz only
	// emits a full analysis (DNSKEY set, DS at parent, queries…) for names
	// passed on the command line: ancestors probed implicitly are kept as
	// "stub" entries that grok ignores. Listing them explicitly is what
	// makes the chain (root, TLD, intermediates, leaf) appear in the grok
	// output, and therefore in the report.
	names := ancestorNames(domain)

	probeOut, probeErr, err := runProbe(probeCtx, bin, names, extraArgs)
	if err != nil {
		return nil, fmt.Errorf("dnsviz probe failed: %w (stderr: %s)", err, truncate(probeErr, 4096))
	}

	grokCtx, cancelGrok := context.WithTimeout(ctx, timeout)
	defer cancelGrok()

	trustAnchors := c.TrustAnchorsFile
	if trustAnchors == "" {
		if _, err := os.Stat(defaultTrustAnchorsFile); err == nil {
			trustAnchors = defaultTrustAnchorsFile
		}
	}

	grokOut, grokErr, err := runGrok(grokCtx, bin, probeOut, trustAnchors)
	if err != nil {
		return nil, fmt.Errorf("dnsviz grok failed: %w (stderr: %s)", err, truncate(grokErr, 4096))
	}

	zones, order, err := checker.ParseGrokOutput(grokOut)
	if err != nil {
		return nil, fmt.Errorf("decoding dnsviz grok output: %w", err)
	}

	return &checker.DNSVizData{
		Domain:      domain,
		Zones:       zones,
		Order:       order,
		Raw:         grokOut,
		ProbeStderr: probeErr,
		GrokStderr:  grokErr,
	}, nil
}

func runProbe(ctx context.Context, bin string, names []string, extraArgs string) ([]byte, string, error) {
	args := []string{"probe"}
	args = append(args, strings.Fields(extraArgs)...)
	args = append(args, "--")
	args = append(args, names...)

	cmd := exec.CommandContext(ctx, bin, args...)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &capLimit{B: &stdout, max: maxDNSVizOutputBytes}
	cmd.Stderr = &capLimit{B: &stderr, max: maxDNSVizOutputBytes}
	if err := cmd.Run(); err != nil {
		return stdout.Bytes(), stderr.String(), err
	}
	return stdout.Bytes(), stderr.String(), nil
}

func runGrok(ctx context.Context, bin string, probeJSON []byte, trustAnchorsFile string) ([]byte, string, error) {
	args := []string{"grok"}
	if trustAnchorsFile != "" {
		args = append(args, "-t", trustAnchorsFile)
	}
	cmd := exec.CommandContext(ctx, bin, args...)
	cmd.Stdin = bytes.NewReader(probeJSON)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &capLimit{B: &stdout, max: maxDNSVizOutputBytes}
	cmd.Stderr = &capLimit{B: &stderr, max: maxDNSVizOutputBytes}
	if err := cmd.Run(); err != nil {
		return stdout.Bytes(), stderr.String(), err
	}
	return stdout.Bytes(), stderr.String(), nil
}

// ancestorNames returns the root, each parent label suffix, and finally
// the queried name itself, all in trailing-dot form. For
// "sub.example.com" the result is [".", "com.", "example.com.",
// "sub.example.com."]. The order matters: dnsviz probes names in the
// order they're given and creates a "stub" entry for each ancestor
// referenced by a previously-seen name, so listing root → leaf is the
// only ordering that yields a fully-analyzed entry for every ancestor
// in the grok output. dnsviz tolerates non-zone names on the command
// line (the analysis attaches to the enclosing zone), so we don't need
// to pre-compute the real zone cuts.
func ancestorNames(domain string) []string {
	domain = strings.TrimSuffix(domain, ".")
	if domain == "" {
		return []string{"."}
	}
	rev := []string{domain + "."}
	for {
		i := strings.Index(domain, ".")
		if i < 0 {
			break
		}
		domain = domain[i+1:]
		rev = append(rev, domain+".")
	}
	rev = append(rev, ".")
	// Reverse to root → leaf.
	out := make([]string, len(rev))
	for i, n := range rev {
		out[len(rev)-1-i] = n
	}
	return out
}

func isValidDomainName(s string) bool {
	if s == "" || len(s) > 253 || s[0] == '-' || s[0] == '.' {
		return false
	}
	for i := 0; i < len(s); i++ {
		c := s[i]
		switch {
		case c >= 'a' && c <= 'z':
		case c >= 'A' && c <= 'Z':
		case c >= '0' && c <= '9':
		case c == '-' || c == '.' || c == '_':
		default:
			return false
		}
	}
	return true
}

// truncate returns s capped at n bytes, trimming back to a UTF-8 rune
// boundary so the appended ellipsis can't follow a half-encoded codepoint.
func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	for n > 0 && !utf8.RuneStart(s[n]) {
		n--
	}
	return s[:n] + "…"
}

// capLimit is an io.Writer that buffers up to max bytes into B and silently
// drops anything beyond the cap. It exists to keep a runaway dnsviz process
// from filling memory while still letting os/exec consume the pipe.
//
// Write deliberately reports len(p) bytes written even when it discarded
// the overflow: returning a "short write" would make os/exec abort the
// child with io.ErrShortWrite, which is not what we want: we want to keep
// reading the rest of the output (and trash it) until the process exits.
// The trade-off is that the returned count lies; callers must not rely on
// it for accounting.
type capLimit struct {
	B   *bytes.Buffer
	max int
}

func (c *capLimit) Write(p []byte) (int, error) {
	remaining := c.max - c.B.Len()
	if remaining <= 0 {
		return len(p), nil
	}
	if len(p) > remaining {
		c.B.Write(p[:remaining])
		return len(p), nil
	}
	return c.B.Write(p)
}