Skip to content

Commit 3f707e1

Browse files
terrorbyteterrorbyte
committed
Add tests, docs, linter fixes, and some simplifications.
1 parent c1c00eb commit 3f707e1

File tree

5 files changed

+444
-149
lines changed

5 files changed

+444
-149
lines changed

‎cli/commandline.go‎

Lines changed: 58 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -837,91 +837,70 @@ func LocalCmdLineParse(conf *config.Config) bool {
837837
return handleLogOptions(logFile, frameworkLogLevel, exploitLogLevel)
838838
}
839839

840-
// Adds default flags for payload types, this allows classes of payloads that are supported to
841-
// use globally defined command line flags without having to redifine them each exploit.
842-
func addPayloadFlags(conf *config.Config) {
843-
if conf.PayloadFlags {
844-
if len(conf.SupportedPayloads) == 1 {
845-
conf.SupportedPayloads[0].Default = payload.Default
846-
}
847-
hasDefault := false
848-
defaultType := ""
849-
defaultArch := ""
850-
typeOptions := []string{}
851-
archOptions := []string{}
852-
count := map[payload.Types]int{}
853-
for i, supported := range conf.SupportedPayloads {
854-
switch supported.Types {
855-
case payload.LinuxCommand:
856-
_, exists := conf.StringFlagsMap["command"]
857-
if !exists {
858-
conf.CreateStringFlag("command", "", "Command to use for the exploit, an empty string will use the exploit default.")
859-
}
860-
case payload.WindowsCommand:
861-
_, exists := conf.StringFlagsMap["command"]
862-
if !exists {
863-
conf.CreateStringFlag("command", "", "Command to use for the exploit, an empty string will use the exploit default.")
864-
}
865-
case payload.MacCommand:
866-
_, exists := conf.StringFlagsMap["command"]
867-
if !exists {
868-
conf.CreateStringFlag("command", "", "Command to use for the exploit, an empty string will use the exploit default.")
869-
}
870-
case payload.GenericCommand:
871-
_, exists := conf.StringFlagsMap["command"]
872-
if !exists {
873-
conf.CreateStringFlag("command", "", "Command to use for the exploit, an empty string will use the exploit default.")
874-
}
875-
case payload.LinuxELF:
876-
_, exists := conf.StringFlagsMap["payload"]
877-
if !exists {
878-
conf.CreateStringFlag("payload", "", "Path to load custom payload from, an empty string will use the exploit default.")
879-
}
880-
case payload.LinuxSO:
881-
_, exists := conf.StringFlagsMap["payload"]
882-
if !exists {
883-
conf.CreateStringFlag("payload", "", "Path to load custom payload from, an empty string will use the exploit default.")
884-
}
885-
case payload.WindowsEXE:
886-
_, exists := conf.StringFlagsMap["payload"]
887-
if !exists {
888-
conf.CreateStringFlag("payload", "", "Path to load custom payload from, an empty string will use the exploit default.")
889-
}
890-
891-
case payload.WindowsDLL:
892-
_, exists := conf.StringFlagsMap["payload"]
893-
if !exists {
894-
conf.CreateStringFlag("payload", "", "Path to load custom payload from, an empty string will use the exploit default.")
895-
}
896-
case payload.Webshell:
897-
_, exists := conf.StringFlagsMap["payload"]
898-
if !exists {
899-
conf.CreateStringFlag("payload", "", "Path to load custom payload from, an empty string will use the exploit default.")
900-
}
901-
default:
902-
output.PrintFrameworkError("Unexpected payload type used")
840+
func addDefaultPayloadFlags(conf *config.Config) (string, string, map[payload.Type]int, []string, []string) {
841+
if len(conf.SupportedPayloads) == 1 {
842+
conf.SupportedPayloads[0].Default = payload.Default
843+
}
844+
hasDefault := false
845+
defaultType := ""
846+
defaultArch := ""
847+
typeOptions := []string{}
848+
archOptions := []string{}
849+
count := map[payload.Type]int{}
850+
for i, supported := range conf.SupportedPayloads {
851+
switch supported.Type {
852+
case payload.LinuxCommand,
853+
payload.WindowsCommand,
854+
payload.WindowsPowerShellCommand,
855+
payload.MacCommand,
856+
payload.GenericCommand:
857+
_, exists := conf.StringFlagsMap["command"]
858+
if !exists {
859+
conf.CreateStringFlag("command", "", "Command to use for the exploit, an empty string will use the exploit default.")
903860
}
861+
case payload.LinuxELF,
862+
payload.LinuxSO,
863+
payload.WindowsEXE,
864+
payload.WindowsDLL,
865+
payload.Webshell:
866+
_, exists := conf.StringFlagsMap["payload"]
867+
if !exists {
868+
conf.CreateStringFlag("payload", "", "Path to load custom payload from, an empty string will use the exploit default.")
869+
}
870+
case payload.UnspecifiedType:
871+
output.PrintFrameworkError("Unspecified payload type used")
872+
default:
873+
output.PrintFrameworkError("Unexpected payload type used")
874+
}
904875

905-
count[supported.Types]++
906-
typeOptions = append(typeOptions, supported.Types.String())
907-
archOptions = append(archOptions, supported.Arch.String())
908-
if i == 0 && len(conf.SupportedPayloads) == 1 {
909-
defaultType = supported.Types.String()
910-
defaultArch = supported.Arch.String()
876+
count[supported.Type]++
877+
typeOptions = append(typeOptions, supported.Type.String())
878+
archOptions = append(archOptions, supported.Arch.String())
879+
if i == 0 && len(conf.SupportedPayloads) == 1 {
880+
defaultType = supported.Type.String()
881+
defaultArch = supported.Arch.String()
911882

912-
continue
913-
}
914-
if hasDefault && supported.Default == payload.Default {
915-
output.PrintfFrameworkWarn("Multiple default payloads selected, using the first and skipping: %s", supported.Types.String())
883+
continue
884+
}
885+
if hasDefault && supported.Default == payload.Default {
886+
output.PrintfFrameworkWarn("Multiple default payloads selected, using the first and skipping: %s", supported.Type.String())
916887

917-
continue
918-
}
919-
if !hasDefault && supported.Default == payload.Default {
920-
defaultType = supported.Types.String()
921-
defaultArch = supported.Arch.String()
922-
}
888+
continue
889+
}
890+
if !hasDefault && supported.Default == payload.Default {
891+
defaultType = supported.Type.String()
892+
defaultArch = supported.Arch.String()
923893
}
894+
}
895+
896+
return defaultType, defaultArch, count, typeOptions, archOptions
897+
}
924898

899+
// Adds default flags for payload types, this allows classes of payloads that are supported to
900+
// use globally defined command line flags without having to redifine them each exploit.
901+
func addPayloadFlags(conf *config.Config) {
902+
if conf.PayloadFlags {
903+
defaultType, defaultArch, count, typeOptions, archOptions := addDefaultPayloadFlags(conf)
925904
if len(conf.SupportedPayloads) > 1 {
926905
if defaultType == "" {
927906
output.PrintFrameworkError("No default payload type was defined.")

‎config/config.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -505,6 +505,18 @@ func (conf *Config) ResolveC2Payload() c2.Impl {
505505
// }
506506
// // continue with non-custom payload generation
507507
// }
508+
//
509+
// Alternatively, simple payloads can utilize the payload type options
510+
// during payload generations to substitute in the custom payloads by using
511+
// the type specific checks:
512+
//
513+
// if conf.HasCustomPayload() {
514+
// if conf.SelectedPayload.Type.IsCommand() {
515+
// output.PrintfStatus("using '%s' in place of default", string(conf.CustomPayload))
516+
// } else {
517+
// output.PrintfStatus("using binary len %d in place of default", len(string(conf.CustomPayload)))
518+
// }
519+
// }
508520
func (conf *Config) AddPayload(p ...payload.Supported) {
509521
conf.PayloadFlags = true
510522
hasDefault := false

‎framework.go‎

Lines changed: 64 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -472,6 +472,69 @@ func updateGoDebug() {
472472
}
473473
}
474474

475+
func addPayloadMetadata(conf *config.Config) {
476+
cmd := conf.StringFlagsMap["command"]
477+
p := conf.StringFlagsMap["payload"]
478+
if p != nil && cmd != nil {
479+
if *p != "" && *cmd != "" {
480+
output.PrintFrameworkError("Both `-command` and `-payload` cannot be set at the same time.")
481+
482+
return
483+
}
484+
}
485+
payloadType := conf.StringFlagsMap["payload-type"]
486+
if payloadType != nil {
487+
p := payload.TypeFromString(*payloadType)
488+
if p == payload.UnspecifiedType {
489+
output.PrintfFrameworkError("payload type not supported: %s", *payloadType)
490+
491+
return
492+
}
493+
for _, t := range conf.SupportedPayloads {
494+
if t.Type == p {
495+
conf.SelectedPayload = t
496+
}
497+
}
498+
}
499+
if cmd != nil {
500+
if *cmd != "" {
501+
conf.CustomPayload = []byte(*cmd)
502+
}
503+
}
504+
if p != nil {
505+
switch conf.SelectedPayload.Type {
506+
case payload.LinuxSO,
507+
payload.LinuxELF,
508+
payload.WindowsEXE,
509+
payload.WindowsDLL,
510+
payload.Webshell:
511+
if *p == "" {
512+
output.PrintError("Selected payload type requires `-payload` to be used")
513+
514+
return
515+
}
516+
case payload.GenericCommand,
517+
payload.WindowsPowerShellCommand,
518+
payload.WindowsCommand,
519+
payload.MacCommand,
520+
payload.LinuxCommand:
521+
case payload.UnspecifiedType:
522+
output.PrintFrameworkError("Unspecified payload type used")
523+
default:
524+
}
525+
if *p != "" {
526+
d, err := os.ReadFile(*p)
527+
if err != nil {
528+
output.PrintfFrameworkError("Could not read custom payload '%s': %s", *p, err.Error())
529+
530+
return
531+
}
532+
conf.CustomPayload = d
533+
}
534+
}
535+
output.PrintfFrameworkDebug("selecting payload type: %s", *payloadType)
536+
}
537+
475538
// Effectively the package main function. Parses configuration, starts command and control,
476539
// controls which targets are scanned, initiates call down into the exploits implementation
477540
// and is ultimately responsible for waiting for all c2 and attack threads to finish.
@@ -487,56 +550,7 @@ func RunProgram(sploit Exploit, conf *config.Config) {
487550
}
488551
// Early start the payload flags
489552
if conf.PayloadFlags {
490-
cmd := conf.StringFlagsMap["command"]
491-
p := conf.StringFlagsMap["payload"]
492-
if p != nil && cmd != nil {
493-
if *p != "" && *cmd != "" {
494-
output.PrintFrameworkError("Both `-command` and `-payload` cannot be set at the same time.")
495-
496-
return
497-
}
498-
}
499-
payloadType := conf.StringFlagsMap["payload-type"]
500-
if payloadType != nil {
501-
p, ok := payload.TypeFromString(*payloadType)
502-
if !ok {
503-
output.PrintfFrameworkError("payload type not supported: %s", *payloadType)
504-
505-
return
506-
}
507-
for _, t := range conf.SupportedPayloads {
508-
if t.Types == p {
509-
conf.SelectedPayload = t
510-
}
511-
}
512-
}
513-
if cmd != nil {
514-
if *cmd != "" {
515-
conf.CustomPayload = []byte(*cmd)
516-
}
517-
}
518-
if p != nil {
519-
switch conf.SelectedPayload.Types {
520-
case payload.LinuxSO, payload.LinuxELF, payload.WindowsEXE, payload.WindowsDLL, payload.Webshell:
521-
if *p == "" {
522-
output.PrintError("Selected payload type requires `-payload` to be used")
523-
524-
return
525-
}
526-
case payload.GenericCommand, payload.WindowsCommand, payload.MacCommand, payload.LinuxCommand:
527-
default:
528-
}
529-
if *p != "" {
530-
d, err := os.ReadFile(*p)
531-
if err != nil {
532-
output.PrintfFrameworkError("Could not read custom payload '%s': %s", *p, err.Error())
533-
534-
return
535-
}
536-
conf.CustomPayload = d
537-
}
538-
}
539-
output.PrintfFrameworkDebug("selecting payload type: %s", *payloadType)
553+
addPayloadMetadata(conf)
540554
}
541555

542556
// create and init the db if the user provided a database

0 commit comments

Comments
 (0)