main_completion.go

package main

import (
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/TMHSDigital/subenum/internal/dns"
	"github.com/TMHSDigital/subenum/internal/output"
)

// Shell completions and the man page are generated from the FlagSet, the
// same source as -h and the site's flag table, so none of them can drift
// (#88). `subenum completion <shell>` and `subenum man` print them;
// release builds ship them via `subenum generate-assets <dir>`.

// fileFlags take a path, so shells should complete file names for them. A
// test fails when a flag whose usage mentions a file is missing here (#117).
var fileFlags = map[string]bool{
	"w": true, "dL": true, "o": true, "stats": true, "inconclusive": true, "takeover-providers": true, "diff": true, "r": true, "exclude-file": true,
	"simulate-zone": true, "seeds": true, "resume": true, "state": true, "sarif": true,
}

// completionShells are the shells `subenum completion` writes scripts for.
var completionShells = []string{"bash", "zsh", "fish", "powershell"}

// valueChoices lists the fixed values some flags accept, from the lists the
// flags are checked against (#264): every record type, then the common
// combinations.
var valueChoices = map[string][]string{
	"format": output.Formats(),
	"type":   append(dns.Types(), "A,AAAA", "A,AAAA,CNAME"),
}

// subcommands are the non-scan words subenum accepts as its first argument.
var subcommands = []string{"completion", "man", "mcp"}

type flagInfo struct {
	name, arg, usage string
	isBool           bool
}

func flagInfos() []flagInfo {
	_, _, fs, _ := parseFlags(nil)
	var out []flagInfo
	fs.VisitAll(func(fl *flag.Flag) {
		arg, usage := flag.UnquoteUsage(fl)
		b, ok := fl.Value.(interface{ IsBoolFlag() bool })
		out = append(out, flagInfo{name: fl.Name, arg: arg, usage: usage, isBool: ok && b.IsBoolFlag()})
	})
	return out
}

// runSubcommand handles `completion`, `man` and the hidden `generate-assets`.
// ok is false when args[0] is not a subcommand.
func runSubcommand(args []string, stdout, stderr io.Writer) (code int, ok bool) {
	if len(args) == 0 {
		return 0, false
	}
	switch args[0] {
	case "completion":
		if len(args) != 2 {
			_, _ = fmt.Fprintln(stderr, "Usage: subenum completion bash|zsh|fish|powershell")
			return exitUsage, true
		}
		if err := writeCompletion(stdout, args[1]); err != nil {
			_, _ = fmt.Fprintf(stderr, "Error: %v\n", err)
			return exitUsage, true
		}
		return exitOK, true
	case "man":
		writeManPage(stdout)
		return exitOK, true
	case "generate-assets":
		if len(args) != 2 {
			_, _ = fmt.Fprintln(stderr, "Usage: subenum generate-assets <dir>")
			return exitUsage, true
		}
		if err := generateAssets(args[1]); err != nil {
			_, _ = fmt.Fprintf(stderr, "Error: %v\n", err)
			return exitFailure, true
		}
		return exitOK, true
	}
	return 0, false
}

// completionFiles maps each shell to the file name packagers expect.
var completionFiles = map[string]string{
	"bash":       "subenum.bash",
	"zsh":        "_subenum",
	"fish":       "subenum.fish",
	"powershell": "subenum.ps1",
}

func writeCompletion(w io.Writer, shell string) error {
	switch shell {
	case "bash":
		writeBash(w)
	case "zsh":
		writeZsh(w)
	case "fish":
		writeFish(w)
	case "powershell":
		writePowerShell(w)
	default:
		return fmt.Errorf("unknown shell %q (want bash, zsh, fish or powershell)", shell)
	}
	return nil
}

func writeBash(w io.Writer) {
	var names, files []string
	for _, f := range flagInfos() {
		names = append(names, "-"+f.name)
		if fileFlags[f.name] {
			files = append(files, "-"+f.name)
		}
	}
	p := func(format string, a ...any) { _, _ = fmt.Fprintf(w, format, a...) }
	p("# bash completion for subenum. Generated by `subenum completion bash`.\n")
	p("_subenum() {\n")
	p("  local cur=\"${COMP_WORDS[COMP_CWORD]}\" prev=\"${COMP_WORDS[COMP_CWORD-1]}\"\n")
	p("  case \"$prev\" in\n")
	p("    %s) COMPREPLY=( $(compgen -f -- \"$cur\") ); return ;;\n", strings.Join(files, "|"))
	p("    completion) [[ $COMP_CWORD -eq 2 ]] && COMPREPLY=( $(compgen -W %q -- \"$cur\") ); return ;;\n", strings.Join(completionShells, " "))
	for _, name := range sortedKeys(valueChoices) {
		p("    -%s) COMPREPLY=( $(compgen -W %q -- \"$cur\") ); return ;;\n", name, strings.Join(valueChoices[name], " "))
	}
	p("  esac\n")
	p("  if [[ \"$cur\" == -* ]]; then\n")
	p("    COMPREPLY=( $(compgen -W %q -- \"$cur\") )\n", strings.Join(names, " "))
	p("  elif [[ $COMP_CWORD -eq 1 ]]; then\n")
	p("    COMPREPLY=( $(compgen -W %q -- \"$cur\") )\n", strings.Join(subcommands, " "))
	p("  fi\n")
	p("}\n")
	p("complete -o default -F _subenum subenum\n")
}

func writeZsh(w io.Writer) {
	esc := strings.NewReplacer(`\`, `\\`, `[`, `\[`, `]`, `\]`, `'`, `'\''`)
	p := func(format string, a ...any) { _, _ = fmt.Fprintf(w, format, a...) }
	p("#compdef subenum\n# zsh completion for subenum. Generated by `subenum completion zsh`.\n\n")
	p("_arguments \\\n")
	for _, f := range flagInfos() {
		spec := "-" + f.name + "[" + esc.Replace(f.usage) + "]"
		switch {
		case f.isBool:
		case fileFlags[f.name]:
			spec += ":file:_files"
		case valueChoices[f.name] != nil:
			spec += ":" + f.arg + ":(" + strings.Join(valueChoices[f.name], " ") + ")"
		default:
			spec += ":" + f.arg + ":"
		}
		p("  '%s' \\\n", spec)
	}
	p("  '1:subcommand or domain:(%s)'\n", strings.Join(subcommands, " "))
}

func writeFish(w io.Writer) {
	esc := strings.NewReplacer(`\`, `\\`, `'`, `\'`)
	p := func(format string, a ...any) { _, _ = fmt.Fprintf(w, format, a...) }
	p("# fish completion for subenum. Generated by `subenum completion fish`.\n")
	p("complete -c subenum -f -n '__fish_use_subcommand' -a '%s'\n", strings.Join(subcommands, " "))
	p("complete -c subenum -f -n '__fish_seen_subcommand_from completion' -a '%s'\n", strings.Join(completionShells, " "))
	for _, f := range flagInfos() {
		line := "complete -c subenum -o " + f.name
		switch {
		case f.isBool:
		case fileFlags[f.name]:
			line += " -r -F"
		case valueChoices[f.name] != nil:
			line += " -x -a '" + strings.Join(valueChoices[f.name], " ") + "'"
		default:
			line += " -x"
		}
		p("%s -d '%s'\n", line, esc.Replace(f.usage))
	}
}

// writePowerShell completes by the previous word: a file flag returns
// nothing so PowerShell falls back to path completion, a flag with fixed
// values offers them, and otherwise flags (and, first, subcommands) are
// offered (#117).
func writePowerShell(w io.Writer) {
	quote := func(ss []string) string {
		q := make([]string, len(ss))
		for i, s := range ss {
			q[i] = "'" + strings.ReplaceAll(s, "'", "''") + "'"
		}
		return "@(" + strings.Join(q, ", ") + ")"
	}
	var flags, files, valued []string
	for _, f := range flagInfos() {
		flags = append(flags, "-"+f.name)
		switch {
		case fileFlags[f.name]:
			files = append(files, "-"+f.name)
		case !f.isBool && valueChoices[f.name] == nil:
			valued = append(valued, "-"+f.name)
		}
	}
	p := func(format string, a ...any) { _, _ = fmt.Fprintf(w, format, a...) }
	p("# PowerShell completion for subenum. Generated by `subenum completion powershell`.\n")
	p("Register-ArgumentCompleter -Native -CommandName subenum -ScriptBlock {\n")
	p("    param($wordToComplete, $commandAst, $cursorPosition)\n")
	p("    $words = @($commandAst.CommandElements | Where-Object { $_.Extent.EndOffset -lt $cursorPosition } | ForEach-Object { $_.ToString() })\n")
	p("    $prev = if ($words.Count -gt 0) { $words[-1] } else { '' }\n")
	p("    $fileFlags = %s\n", quote(files))
	p("    $valueFlags = %s\n", quote(valued))
	p("    if ($fileFlags -contains $prev -or $valueFlags -contains $prev) { return }\n")
	p("    $values = $null\n")
	p("    switch ($prev) {\n")
	for _, name := range sortedKeys(valueChoices) {
		p("        '-%s' { $values = %s }\n", name, quote(valueChoices[name]))
	}
	p("        'completion' { if ($words.Count -eq 2) { $values = %s } }\n", quote(completionShells))
	p("    }\n")
	p("    if ($null -ne $values) {\n")
	p("        $values | Where-Object { $_ -like \"$wordToComplete*\" } | ForEach-Object {\n")
	p("            [System.Management.Automation.CompletionResult]::new($_, $_, 'ParameterValue', $_)\n")
	p("        }\n")
	p("        return\n")
	p("    }\n")
	p("    $candidates = %s\n", quote(flags))
	p("    if ($words.Count -le 1) { $candidates += %s }\n", quote(subcommands))
	p("    $candidates | Where-Object { $_ -like \"$wordToComplete*\" } | ForEach-Object {\n")
	p("        [System.Management.Automation.CompletionResult]::new($_, $_, 'ParameterName', $_)\n")
	p("    }\n")
	p("}\n")
}

// writeManPage renders subenum(1) in roff.
func writeManPage(w io.Writer) {
	roff := strings.NewReplacer(`\`, `\e`, "-", `\-`)
	line := func(s string) string {
		s = roff.Replace(s)
		if strings.HasPrefix(s, ".") || strings.HasPrefix(s, "'") {
			s = `\&` + s
		}
		return s
	}
	p := func(format string, a ...any) { _, _ = fmt.Fprintf(w, format, a...) }
	p(".TH SUBENUM 1 \"\" \"%s\" \"User Commands\"\n", formatVersion())
	p(".SH NAME\nsubenum \\- subdomain enumeration you can trust\n")
	p(".SH SYNOPSIS\n.B subenum\n[\\fIoptions\\fR] \\fIdomain\\fR\n.br\n.B subenum\n[\\fIoptions\\fR] \\fB\\-dL\\fR \\fIdomains_file\\fR\n")
	p(".br\n.B subenum completion\n\\fBbash\\fR|\\fBzsh\\fR|\\fBfish\\fR|\\fBpowershell\\fR\n.br\n.B subenum man\n")
	p(".br\n.B subenum mcp\n[\\fB\\-allow\\fR \\fIdomains\\fR] [\\fIflags\\fR]\n")
	p(".SH DESCRIPTION\n%s\n.PP\n%s\n",
		line("subenum brute-forces subdomains by resolving a wordlist against a target domain with a concurrent worker pool. Resolved names go to stdout; progress, diagnostics and a per-query breakdown go to stderr. Without -w it uses a built-in list of the 5,000 most common subdomain labels."),
		line("Only scan domains you own or have explicit written permission to test."))
	p(".SH OPTIONS\n")
	for _, f := range flagInfos() {
		if f.arg != "" {
			p(".TP\n.BI \\-%s \" %s\"\n%s\n", roff.Replace(f.name), f.arg, line(f.usage))
		} else {
			p(".TP\n.B \\-%s\n%s\n", roff.Replace(f.name), line(f.usage))
		}
	}
	p(".SH ENVIRONMENT\n")
	p(".TP\n.B SUBENUM_\\fIFLAG\\fR\n%s\n", line("Default for a flag, with the name upper-cased and dashes turned into underscores: SUBENUM_DNS_SERVER, SUBENUM_RATE, SUBENUM_T. Overrides the config file; a flag on the command line overrides both. Neither applies to subenum mcp, which takes its own flags only."))
	p(".TP\n.B SUBENUM_CONFIG\n%s\n", line("Path of the config file, instead of the default location."))
	p(".SH FILES\n.TP\n.I <user config dir>/subenum/config.json\n%s\n", line(`Default flag values as a JSON object keyed by flag name, e.g. {"dns-server": "1.1.1.1:53", "rate": 200}. Run subenum -print-config to see the effective settings and their sources.`))
	p(".SH EXIT STATUS\n")
	for _, l := range strings.Split(strings.TrimSpace(exitCodesHelp), "\n")[1:] {
		fields := strings.Fields(l)
		p(".TP\n.B %s\n%s\n", fields[0], line(strings.Join(fields[1:], " ")))
	}
	p(".SH SEE ALSO\nhttps://tmhsdigital.github.io/subenum/\n")
}

// generateAssets writes the completions and the man page under dir, for
// release archives.
func generateAssets(dir string) error {
	for _, sub := range []string{"completions", "man"} {
		if err := os.MkdirAll(filepath.Join(dir, sub), 0o755); err != nil { //nolint:gosec // build output
			return err
		}
	}
	write := func(path string, render func(io.Writer)) error {
		f, err := os.Create(path) //nolint:gosec // path comes from the release build, not user input
		if err != nil {
			return err
		}
		render(f)
		return f.Close()
	}
	for _, shell := range sortedKeys(completionFiles) {
		path := filepath.Join(dir, "completions", completionFiles[shell])
		if err := write(path, func(w io.Writer) { _ = writeCompletion(w, shell) }); err != nil {
			return err
		}
	}
	return write(filepath.Join(dir, "man", "subenum.1"), writeManPage)
}

func sortedKeys[V any](m map[string]V) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	return keys
}