Skip to content

[WIP] Improve help message (grouping) #760

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
+195 −2
Draft

[WIP] Improve help message (grouping) #760

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
197 changes: 195 additions & 2 deletions cmd/root.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,26 @@ import (
"os"
"path/filepath"
"runtime"
"strings"
"text/template"
"time"
"unicode"
"unicode/utf8"

"github.com/hahwul/dalfox/v2/internal/har"
"github.com/hahwul/dalfox/v2/internal/printing"
"github.com/hahwul/dalfox/v2/pkg/model"
"github.com/logrusorgru/aurora"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
)

// FlagGroup is a struct for grouping flags
type FlagGroup struct {
Title string
Flags *pflag.FlagSet
}

// Default option values
const (
DefaultCustomAlertValue = "1"
Expand All @@ -33,12 +44,73 @@ const (
var options model.Options
var harFilePath string
var args Args
var flagGroups []FlagGroup

const customHelpTemplate = `{{.LongOrUsage}}

Usage:
{{.Command.UseLine}}{{if .Command.HasAvailableSubCommands}} [command]{{end}}
{{if .Command.HasAvailableSubCommands}}
Available Commands:{{range .Command.Commands}}{{if (or .IsAvailableCommand .IsAdditionalHelpTopicCommand)}}
{{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}
{{end}}
{{if .ShowFlagGroups }}
{{range .FlagGroupsRef}}
{{.Title}}:
{{.Flags.FlagUsages | trimTrailingWhitespaces}}
{{end}}
{{end}}
{{if .Command.HasAvailableLocalFlags}}
{{if .Command.HasParent}} {{/* For subcommands, this is fine */}}
Local Flags:
{{.Command.LocalFlags.FlagUsages | trimTrailingWhitespaces}}
{{else if .ShowFlagGroups}} {{/* For root command with flag groups shown */}}
{{ $helpFlag := .Command.LocalFlags.Lookup "help" }}
{{ if $helpFlag }}
Local Flags:
-h, --help {{$helpFlag.Usage}}
{{end}}
Comment on lines +67 to +72
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The current template logic for the root command, when flag groups are shown (.ShowFlagGroups is true), will only display the -h, --help flag under "Local Flags". If rootCmd were to have other local (non-persistent) flags defined directly on it, they would not be visible in the help output when groups are active.

Is this the intended behavior? If rootCmd might have other local flags, consider adjusting the template to ensure they are always displayed, perhaps in their own distinct section if not part of the --help flag logic.

{{else}} {{/* For root command if flag groups are NOT shown (e.g. vanilla help) */}}
Local Flags:
{{.Command.LocalFlags.FlagUsages | trimTrailingWhitespaces}}
{{end}}
{{end}}
{{if .Command.HasAvailableInheritedFlags}}{{if not .ShowFlagGroups }}
Global Flags:
{{.Command.InheritedFlags.FlagUsages | trimTrailingWhitespaces}}
{{end}}{{end}}
{{if .Command.HasExample}}
Examples:
{{.Command.Example}}
{{end}}
{{if .Command.HasHelpSubCommands}}
Additional help topics:{{range .Command.Commands}}{{if .IsAdditionalHelpTopicCommand}}
{{rpad .Command.CommandPath .Command.CommandPathPadding}} {{.Short}}{{end}}{{end}}
{{end}}
{{if .Command.HasAvailableSubCommands}}
Use "{{.Command.CommandPath}} [command] --help" for more information about a command.{{end}}
`

var rootCmd = &cobra.Command{
Use: "dalfox",
Short: "Dalfox is a powerful open-source XSS scanner and utility focused on automation.",
Long: `Dalfox is a fast and powerful parameter analysis and XSS scanning tool.
It helps you find XSS vulnerabilities in web applications with ease.
Dalfox supports various features like parameter mining, custom payloads,
blind XSS detection, and much more.`,
Run: func(cmd *cobra.Command, args []string) {
printing.Banner(options)
printing.DalLog("YELLOW", "Read the help page using the -h flag to see other options and flags!", options)
// If no subcommand is given and it's not a help request, show help.
// This prevents running the banner and log if just 'dalfox' is typed.
// Help flag (-h, --help) is handled by Cobra automatically by showing help.
if len(args) == 0 {
cmd.Help()
os.Exit(0)
}
Comment on lines 101 to +108
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The Run function for rootCmd now shows help and exits if no arguments are provided. This is a common and user-friendly behavior.

However, the previous behavior of showing printing.Banner(options) (lines 112-113, now commented out) when dalfox is run without arguments will no longer occur. If the banner should still be displayed before the help message, you might consider calling printing.Banner(options) before cmd.Help().

// Original run logic (perhaps for a default action if args were present but not a known command)
// For now, if args are present and not a command, Cobra shows an error.
// If you want 'dalfox somearg' to do something specific by default, that logic goes here.
// printing.Banner(options)
// printing.DalLog("YELLOW", "Read the help page using the -h flag to see other options and flags!", options)
},
}

Expand Down Expand Up @@ -123,6 +195,127 @@ func init() {
rootCmd.PersistentFlags().BoolVar(&args.OutputResponse, "output-response", false, "Include raw HTTP responses in the results. Example: --output-response")
rootCmd.PersistentFlags().BoolVar(&args.SkipDiscovery, "skip-discovery", false, "Skip the entire discovery phase, proceeding directly to XSS scanning. Requires -p flag to specify parameters. Example: --skip-discovery -p 'username'")
rootCmd.PersistentFlags().BoolVar(&args.ForceHeadlessVerification, "force-headless-verification", false, "Force headless browser-based verification, useful when automatic detection fails or to override default behavior. Example: --force-headless-verification")

// Initialize flag groups
initializeFlagGroups()

// Set custom help template and function for rootCmd
rootCmd.SetHelpTemplate(customHelpTemplate)
rootCmd.SetHelpFunc(getCustomHelpFunction())

// Apply to subcommands
for _, subCmd := range rootCmd.Commands() {
subCmd.SetHelpTemplate(customHelpTemplate)
subCmd.SetHelpFunc(getCustomHelpFunction()) // Use the same custom help logic
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The getCustomHelpFunction() is called for rootCmd and then again for each subcommand inside the loop. This creates a new, identical closure for the help function each time. If the help function's logic doesn't need to capture different states per command (beyond the command argument passed to it by Cobra, which it doesn't seem to), you could define the help function once and reuse that instance. This would be a minor optimization and could slightly improve clarity.

For example, in init():

// Define the help function once
customHelpFn := func(command *cobra.Command, args []string) {
    // ... (current logic from getCustomHelpFunction's returned func)
}

// Assign it
rootCmd.SetHelpFunc(customHelpFn)
for _, subCmd := range rootCmd.Commands() {
    subCmd.SetHelpFunc(customHelpFn)
}

This avoids the overhead of the getCustomHelpFunction call and closure creation for each command.

Suggested change
subCmd.SetHelpFunc(getCustomHelpFunction()) // Use the same custom help logic
subCmd.SetHelpFunc(rootCmd.HelpFunc()) // Use the same custom help logic, potentially reusing the root command's already set function if it's identical.

}
}

// getCustomHelpFunction returns a function that can be used as a custom help function for cobra commands.
// This approach allows us to generate a new closure for each command if needed, or just share the same one.
func getCustomHelpFunction() func(*cobra.Command, []string) {
return func(command *cobra.Command, args []string) {
// Data to pass to the template
// The template expects fields like .Command, .flagGroupsRef, .showFlagGroups
templateData := struct {
Command *cobra.Command
FlagGroupsRef []FlagGroup // Renamed to match template expectation if it was .flagGroupsRef
ShowFlagGroups bool
// Expose other necessary fields/methods if your template uses them directly,
// e.g. if it doesn't use .Command.Long but just .Long
LongOrUsage string
}{
Command: command,
FlagGroupsRef: flagGroups, // Assumes flagGroups is accessible (e.g. package-level var)
ShowFlagGroups: len(flagGroups) > 0, // Show groups if they exist
}

// Logic for LongOrUsage (simplified from Cobra's internal help command)
if command.Long != "" {
templateData.LongOrUsage = command.Long
} else if command.UsageString() != "" {
templateData.LongOrUsage = command.UsageString()
} else {
templateData.LongOrUsage = command.Short
}


tmpl := template.New("customHelp")

// Manually add rpad and other necessary functions
tmpl.Funcs(template.FuncMap{
"rpad": func(s string, padding int) string {
// A simple rpad, Cobra's might handle multi-byte chars better by default.
// This one calculates padding based on rune count.
sLen := utf8.RuneCountInString(s)
if padding <= sLen {
return s
}
return s + strings.Repeat(" ", padding-sLen)
},
"trimTrailingWhitespaces": func(s string) string {
return strings.TrimRightFunc(s, unicode.IsSpace)
},
// Cobra's default func map also includes:
// "gt", "hasPrefix", "hasSuffix", "contains", "eq", "ne",
// "split", "replace", "join", "lower", "upper", "title",
// "trim", "trimLeft", "trimRight", "substring"
// Many of these are from text/template or a common library like sprig,
// but rpad is specific.
})

// customHelpTemplate is the global const string
parsedTmpl, err := tmpl.Parse(customHelpTemplate)
if err != nil {
command.PrintErrln("Error parsing custom help template:", err)
command.PrintErrln(command.UsageString())
return
}

err = parsedTmpl.Execute(command.OutOrStdout(), templateData)
if err != nil {
command.PrintErrln("Error executing custom help template:", err)
command.PrintErrln(command.UsageString())
}
}
}


func initializeFlagGroups() {
flagGroups = []FlagGroup{
{Title: "Input", Flags: pflag.NewFlagSet("Input", pflag.ExitOnError)},
{Title: "Request", Flags: pflag.NewFlagSet("Request", pflag.ExitOnError)},
{Title: "Scanning", Flags: pflag.NewFlagSet("Scanning", pflag.ExitOnError)},
{Title: "Mining", Flags: pflag.NewFlagSet("Mining", pflag.ExitOnError)},
{Title: "Output", Flags: pflag.NewFlagSet("Output", pflag.ExitOnError)},
{Title: "Advanced", Flags: pflag.NewFlagSet("Advanced", pflag.ExitOnError)},
}

flagMap := map[string][]string{
"Input": {"config", "custom-payload", "custom-blind-xss-payload", "data", "grep", "remote-payloads", "remote-wordlists", "har-file-path"},
"Request": {"header", "cookie", "user-agent", "method", "cookie-from-raw"},
"Scanning": {"param", "ignore-param", "blind", "timeout", "delay", "worker", "skip-headless", "deep-domxss", "waf-evasion", "skip-discovery", "force-headless-verification", "use-bav", "skip-bav", "skip-mining-dom", "skip-mining-dict", "skip-mining-all", "skip-xss-scanning", "only-custom-payload", "skip-grepping"},
"Mining": {"mining-dict-word", "mining-dict", "mining-dom"},
"Output": {"output", "format", "only-poc", "report", "output-all", "output-request", "output-response", "poc-type", "report-format", "silence", "no-color", "no-spinner"},
"Advanced": {"custom-alert-value", "custom-alert-type", "found-action", "found-action-shell", "proxy", "ignore-return", "max-cpu", "only-discovery", "follow-redirects", "debug"},
}
Comment on lines +293 to +300
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The flagMap defines the association between flag names and their groups. This map is crucial for the correct functioning of the grouped help.

To improve maintainability and reduce the risk of errors (e.g., typos in flag names, or forgetting to add a new flag):

  1. Ensure all existing persistent flags are correctly listed in this map. A small unit test could be beneficial to verify that every persistent flag defined on rootCmd is present in one of the flagMap lists (or handled by an "Other Flags" group as suggested in another comment).
  2. The keys of flagMap (e.g., "Input", "Request") are string literals that must match the Title field of the FlagGroup structs. While simple, this introduces a small chance of inconsistency if a title is updated in one place but not the other. Using constants for group titles or structuring the data to avoid this duplication could be considered for long-term maintainability, though the current approach is acceptable if managed carefully.


rootCmd.PersistentFlags().VisitAll(func(f *pflag.Flag) {
assigned := false
for i, group := range flagGroups {
for _, flagName := range flagMap[group.Title] {
if f.Name == flagName {
flagGroups[i].Flags.AddFlag(f)
assigned = true
break
}
}
if assigned {
break
}
}
// If a flag is not assigned to any group, add it to a default "Other" group or handle as needed.
// For now, we assume all flags will be assigned.
})
Comment on lines +302 to +318
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The current logic in initializeFlagGroups assigns flags to groups based on the flagMap. The comment on line 316-317 (// If a flag is not assigned to any group... For now, we assume all flags will be assigned.) highlights a potential issue: if a new persistent flag is added to rootCmd but is forgotten in flagMap, it will not be displayed in any of the defined groups. Since persistent flags are not typically shown under "Local Flags" when groups are active (unless it's the help flag, per template lines 68-72), such an unmapped flag could become effectively hidden from the help message.

To make this more robust and prevent flags from being accidentally hidden, consider implementing a fallback mechanism, such as an "Other Flags" or "Miscellaneous Flags" group.

One way to do this:

  1. After populating the defined groups, iterate through rootCmd.PersistentFlags().VisitAll() again.
  2. For each flag, check if it has already been added to one of the existing flagGroups.
  3. If a flag hasn't been assigned, add it to a dedicated "Other Flags" pflag.FlagSet.
  4. Append this "Other Flags" group to the flagGroups slice so it gets rendered by the template.

This would ensure all persistent flags are always visible in the help output, improving maintainability.

}

// initConfig reads in config file and ENV variables if set.
Expand Down
Loading