flags.go

package protocol

import (
	"flag"

	"github.com/vulncheck-oss/go-exploit/output"
)

// All arguments used by flag packages use this structure.
type flagArgs struct {
	Variable    any
	Default     any
	Description string
}

// The internalProtocols datastructure is a map based on protocols that then utilizes the interface
// versions of the flag arguments, which should make it generic enough for any of the flag package
// function types that we utilize. This also implies that the protocol variables that are set are
// utilized elsewhere in go-exploit and not freeform CreateStringFlag style flags. If you want to
// add more protocol flags this is where they go.
var internalProtocols = map[string]map[string]flagArgs{
	"HTTP": {
		// allow the user to use their own HTTP user-agent if they so wish
		"user-agent": flagArgs{
			Variable:    GlobalUA,
			Default:     GlobalUA,
			Description: "The User-Agent to use in HTTP requests",
		},
	},
}

// AddProtocolFlags uses the internal supported protocol variables and manually adds the protocol
// flags to the exploit. If your exploit uses multiple protocols and you would like to expose the
// protocol flags, it should be added with a `protocol.AddProtocolFlags("HTTP").
func AddProtocolFlags(protocol string) bool {
	_, exists := internalProtocols[protocol]
	if !exists {
		output.PrintfFrameworkError("Protocol '%s' has no flags specified", protocol)

		return false
	}

	return CreateProtocolFlags(protocol)
}

// CreateProtocolFlags takes a protocol string and matches it to the internalProtocols map and then
// uses the type assertions to make them match the flag package functions. Be sure to check that
// the internalProtocols types for Variable and Default match the flag function signatures or this
// will bail out before the flag setting silently to basic type errors.
func CreateProtocolFlags(protocol string) bool {
	flags, exists := internalProtocols[protocol]
	if !exists {
		// No flags for this protocol enabled, this is not fatal, it just means that no
		// flags have set for the protocol
		return true
	}
	for name, args := range flags {
		switch v := args.Default.(type) {
		case int:
			a, ok := args.Variable.(int)
			if !ok {
				output.PrintFrameworkError("Internal protocol variable was not the expected type (int)")

				return false
			}
			flag.IntVar(&a, name, v, args.Description)
		case uint:
			a, ok := args.Variable.(uint)
			if !ok {
				output.PrintFrameworkError("Internal protocol variable was not the expected type (uint)")

				return false
			}
			flag.UintVar(&a, name, v, args.Description)
		case string:
			a, ok := args.Variable.(string)
			if !ok {
				output.PrintFrameworkError("Internal protocol variable was not the expected type (string)")

				return false
			}
			flag.StringVar(&a, name, v, args.Description)
		case bool:
			a, ok := args.Variable.(bool)
			if !ok {
				output.PrintFrameworkError("Internal protocol variable was not the expected type (bool)")

				return false
			}
			flag.BoolVar(&a, name, v, args.Description)
		default:
			// Type not found, how did this happen?
		}
	}

	return true
}