Add payload support definition & support for custom user payloads (#459)

Author: terrorbyte

c2/channel/channel.go

@@ -80,7 +80,7 @@ func (c *Channel) AddSession(conn *net.Conn, addr string) bool {
 	return true
 }
 
-// Updates the LastSeen value for provided connection to the provided time
+// Updates the LastSeen value for provided connection to the provided time.
 func (c *Channel) UpdateLastSeenByConn(conn net.Conn, timeStamp time.Time) bool {
 	id, ok := c.GetSessionIDByConn(conn)
 	if !ok {
@@ -100,7 +100,7 @@ func (c *Channel) UpdateLastSeenByConn(conn net.Conn, timeStamp time.Time) bool
 	return true
 }
 
-// Returns the session ID that contains a given connection
+// Returns the session ID that contains a given connection.
 func (c *Channel) GetSessionIDByConn(conn net.Conn) (string, bool) {
 	if len(c.Sessions) == 0 {
 		output.PrintFrameworkDebug("No sessions exist")
@@ -119,7 +119,6 @@ func (c *Channel) GetSessionIDByConn(conn net.Conn) (string, bool) {
 	return "", false
 }
 
-
 // RemoveSession removes a specific session ID and if a connection exists, closes it.
 func (c *Channel) RemoveSession(id string) bool {
 	if len(c.Sessions) == 0 {

cli/commandline.go

@@ -15,6 +15,7 @@ import (
 	"github.com/vulncheck-oss/go-exploit/config"
 	"github.com/vulncheck-oss/go-exploit/db"
 	"github.com/vulncheck-oss/go-exploit/output"
+	"github.com/vulncheck-oss/go-exploit/payload"
 	"github.com/vulncheck-oss/go-exploit/protocol"
 )
 
@@ -482,7 +483,10 @@ func printDetails(conf *config.Config) {
 	for _, value := range conf.SupportedC2 {
 		supportedC2Strings = append(supportedC2Strings, value.Name)
 	}
-
+	supportedPayloadsStrings := make([]string, 0)
+	for _, value := range conf.SupportedPayloads {
+		supportedPayloadsStrings = append(supportedPayloadsStrings, value.String())
+	}
 	customFlags := make([]CustomFlag, 0)
 	for key, value := range conf.StringFlagsMap {
 		customFlags = append(customFlags, CustomFlag{
@@ -519,6 +523,7 @@ func printDetails(conf *config.Config) {
 		"VersionScanner", conf.Impl.VersionScanning,
 		"Exploitation", conf.Impl.Exploitation,
 		"SupportedC2", supportedC2Strings,
+		"SupportedPayloads", supportedPayloadsStrings,
 		"Vendor", conf.Vendor,
 		"Products", conf.Products,
 		"CPE", conf.CPE,
@@ -548,6 +553,7 @@ func CodeExecutionCmdLineParse(conf *config.Config) bool {
 	exploitFunctionality(conf)
 	sslFlags(conf)
 	c2Flags(&c2Selection, conf)
+	addPayloadFlags(conf)
 	detailsFlag := flag.Bool("details", false, "Print the implementation details for this exploit")
 
 	flag.Usage = func() {
@@ -612,6 +618,7 @@ func InformationDisclosureCmdLineParse(conf *config.Config) bool {
 	localHostFlags(conf)
 	exploitFunctionality(conf)
 	sslFlags(conf)
+	addPayloadFlags(conf)
 	detailsFlag := flag.Bool("details", false, "Print the implementation details for this exploit")
 
 	flag.Usage = func() {
@@ -654,6 +661,7 @@ func WebShellCmdLineParse(conf *config.Config) bool {
 	localHostFlags(conf)
 	exploitFunctionality(conf)
 	sslFlags(conf)
+	addPayloadFlags(conf)
 	detailsFlag := flag.Bool("details", false, "Print the implementation details for this exploit")
 
 	flag.Usage = func() {
@@ -726,6 +734,7 @@ func FormatFileCmdLineParse(conf *config.Config) bool {
 	localHostFlags(conf)
 	exploitFunctionality(conf)
 	c2Flags(&c2Selection, conf)
+	addPayloadFlags(conf)
 	detailsFlag := flag.Bool("details", false, "Print the implementation details for this exploit")
 	flag.StringVar(&templateFile, "in", "", "The file format template to work with")
 	flag.StringVar(&conf.FileFormatFilePath, "out", "", "The file to write the malicious file to")
@@ -792,6 +801,7 @@ func LocalCmdLineParse(conf *config.Config) bool {
 	localHostFlags(conf)
 	exploitFunctionality(conf)
 	c2Flags(&c2Selection, conf)
+	addPayloadFlags(conf)
 	detailsFlag := flag.Bool("details", false, "Print the implementation details for this exploit")
 
 	flag.Usage = func() {
@@ -826,3 +836,83 @@ func LocalCmdLineParse(conf *config.Config) bool {
 
 	return handleLogOptions(logFile, frameworkLogLevel, exploitLogLevel)
 }
+
+func addDefaultPayloadFlags(conf *config.Config) (string, string, map[payload.Type]int, []string, []string) {
+	if len(conf.SupportedPayloads) == 1 {
+		conf.SupportedPayloads[0].Default = payload.Default
+	}
+	hasDefault := false
+	defaultType := ""
+	defaultArch := ""
+	typeOptions := []string{}
+	archOptions := []string{}
+	count := map[payload.Type]int{}
+	for i, supported := range conf.SupportedPayloads {
+		switch supported.Type {
+		case payload.LinuxCommand,
+			payload.WindowsCommand,
+			payload.WindowsPowerShellCommand,
+			payload.MacCommand,
+			payload.GenericCommand:
+			_, exists := conf.StringFlagsMap["command"]
+			if !exists {
+				conf.CreateStringFlag("command", "", "Command to use for the exploit, an empty string will use the exploit default.")
+			}
+		case payload.LinuxELF,
+			payload.LinuxSO,
+			payload.WindowsEXE,
+			payload.WindowsDLL,
+			payload.Webshell:
+			_, exists := conf.StringFlagsMap["payload"]
+			if !exists {
+				conf.CreateStringFlag("payload", "", "Path to load custom payload from, an empty string will use the exploit default.")
+			}
+		case payload.UnspecifiedType:
+			output.PrintFrameworkError("Unspecified payload type used")
+		default:
+			output.PrintFrameworkError("Unexpected payload type used")
+		}
+
+		count[supported.Type]++
+		typeOptions = append(typeOptions, supported.Type.String())
+		archOptions = append(archOptions, supported.Arch.String())
+		if i == 0 && len(conf.SupportedPayloads) == 1 {
+			defaultType = supported.Type.String()
+			defaultArch = supported.Arch.String()
+
+			continue
+		}
+		if hasDefault && supported.Default == payload.Default {
+			output.PrintfFrameworkWarn("Multiple default payloads selected, using the first and skipping: %s", supported.Type.String())
+
+			continue
+		}
+		if !hasDefault && supported.Default == payload.Default {
+			defaultType = supported.Type.String()
+			defaultArch = supported.Arch.String()
+		}
+	}
+
+	return defaultType, defaultArch, count, typeOptions, archOptions
+}
+
+// Adds default flags for payload types, this allows classes of payloads that are supported to
+// use globally defined command line flags without having to redefine them each exploit.
+func addPayloadFlags(conf *config.Config) {
+	if conf.PayloadFlags {
+		defaultType, defaultArch, count, typeOptions, archOptions := addDefaultPayloadFlags(conf)
+		if len(conf.SupportedPayloads) > 1 {
+			if defaultType == "" {
+				output.PrintFrameworkError("No default payload type was defined.")
+			}
+			conf.CreateStringFlag("payload-type", defaultType, "Payload type to use based on supported types: "+strings.Join(typeOptions, ", "))
+			for _, v := range count {
+				if v > 1 {
+					conf.CreateStringFlag("payload-arch", defaultArch, "Payload architecture to use based on supported archs: "+strings.Join(archOptions, ", "))
+
+					break
+				}
+			}
+		}
+	}
+}