Add protocol based flag handling

cli/commandline.go

@@ -271,8 +271,10 @@ func handleRhostsOptions(conf *config.Config, rhosts string, rports string, rhos
 		return parseRhostsFile(conf, rhostsFile)
 	}
 
-	// dump the selected user agent to debug
-	output.PrintfFrameworkDebug("Using the HTTP User-Agent: %s", protocol.GlobalUA)
+	if conf.Protocol == "HTTP" {
+		// dump the selected user agent to debug
+		output.PrintfFrameworkDebug("Using the HTTP User-Agent: %s", protocol.GlobalUA)
+	}
 
 	return buildRhosts(conf, rhosts, rports)
 }
@@ -387,8 +389,6 @@ func remoteHostFlags(conf *config.Config, rhosts *string, rhostsFile *string, rp
 	flag.StringVar(rports, "rports", "", "A comma delimited list of the remote target's server ports")
 	// generic all comms connection timeout
 	flag.IntVar(&protocol.GlobalCommTimeout, "timeout", 10, "A default timeout for all socket communication")
-	// allow the user to use their own HTTP user-agent if they so wish
-	flag.StringVar(&protocol.GlobalUA, "user-agent", protocol.GlobalUA, "The User-Agent to use in HTTP requests")
 }
 
 // command line flags for defining the local host.
@@ -542,6 +542,11 @@ func CodeExecutionCmdLineParse(conf *config.Config) bool {
 
 	dbFlags(conf)
 	proxyFlags(&proxy)
+	if !protocol.CreateProtocolFlags(conf.Protocol) {
+		output.PrintFrameworkError("internal flag setting for protocols failed")
+
+		return false
+	}
 	loggingFlags(&logFile, &frameworkLogLevel, &exploitLogLevel)
 	remoteHostFlags(conf, &rhosts, &rhostsFile, &rports)
 	localHostFlags(conf)
@@ -607,6 +612,11 @@ func InformationDisclosureCmdLineParse(conf *config.Config) bool {
 
 	dbFlags(conf)
 	proxyFlags(&proxy)
+	if !protocol.CreateProtocolFlags(conf.Protocol) {
+		output.PrintfFrameworkError("internal flag setting for protocols failed")
+
+		return false
+	}
 	loggingFlags(&logFile, &frameworkLogLevel, &exploitLogLevel)
 	remoteHostFlags(conf, &rhosts, &rhostsFile, &rports)
 	localHostFlags(conf)
@@ -649,6 +659,11 @@ func WebShellCmdLineParse(conf *config.Config) bool {
 
 	dbFlags(conf)
 	proxyFlags(&proxy)
+	if !protocol.CreateProtocolFlags(conf.Protocol) {
+		output.PrintfFrameworkError("internal flag setting for protocols failed")
+
+		return false
+	}
 	loggingFlags(&logFile, &frameworkLogLevel, &exploitLogLevel)
 	remoteHostFlags(conf, &rhosts, &rhostsFile, &rports)
 	localHostFlags(conf)

protocol/flags.go

@@ -0,0 +1,97 @@
+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
+}