Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
1 change: 1 addition & 0 deletions Payload_Type/forge/forge/agentfunctions/builder.go
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,7 @@ type collectionSourceCommandData struct {
RepoURL string `json:"repo_url"`
CustomDownloadURL string `json:"custom_download_url"`
CustomVersion string `json:"custom_version"`
ExecutionMethod string `json:"execution_method"` // optional: "execute_assembly" or "inline_assembly"
customAssemblyFileID string
customBofFileIDs []string
customBofExtensionFileID string
Expand Down
6 changes: 5 additions & 1 deletion Payload_Type/forge/forge/agentfunctions/forge_download.go
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,10 @@ func createAssemblyCommand(commandSource collectionSourceCommandData, collection
if commandSource.CustomVersion != "" {
defaultChoices = []string{commandSource.CustomVersion}
}
defaultExecutionMethod := "execute_assembly"
if commandSource.ExecutionMethod != "" {
defaultExecutionMethod = commandSource.ExecutionMethod
}
newCommand := agentstructs.Command{
Name: fmt.Sprintf("%s%s", AssemblyPrefix, commandSource.CommandName),
Description: fmt.Sprintf("%s\nFrom: %s", commandSource.Description, originatingSource),
Expand Down Expand Up @@ -288,7 +292,7 @@ func createAssemblyCommand(commandSource collectionSourceCommandData, collection
ParameterType: agentstructs.COMMAND_PARAMETER_TYPE_CHOOSE_ONE,
Choices: []string{"inline_assembly", "execute_assembly"},
Description: "Specify how the assembly should execute. Execute_assembly is a fork-and-run style architecture, inline_assembly is within the current process.",
DefaultValue: "execute_assembly",
DefaultValue: defaultExecutionMethod,
ModalDisplayName: "Execution Options",
ParameterGroupInformation: []agentstructs.ParameterGroupInfo{
{
Expand Down
29 changes: 28 additions & 1 deletion Payload_Type/forge/forge/agentfunctions/forge_register.go
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,19 @@ func init() {
},
},
},
{
Name: "executionMethod",
ParameterType: agentstructs.COMMAND_PARAMETER_TYPE_CHOOSE_ONE,
Choices: []string{"execute_assembly", "inline_assembly"},
Description: "Override the default execution method for assembly commands",
ModalDisplayName: "Execution Method",
DefaultValue: "",
ParameterGroupInformation: []agentstructs.ParameterGroupInfo{
{
ParameterIsRequired: false,
},
},
},
},
TaskFunctionCreateTasking: func(taskData *agentstructs.PTTaskMessageAllData) agentstructs.PTTaskCreateTaskingMessageResponse {
response := agentstructs.PTTaskCreateTaskingMessageResponse{
Expand All @@ -151,7 +164,14 @@ func init() {
}
remove, err := taskData.Args.GetBooleanArg("remove")
if err != nil {
logging.LogError(err, "failed to get commandName")
logging.LogError(err, "failed to get remove")
response.Success = false
response.Error = err.Error()
return response
}
executionMethod, err := taskData.Args.GetChooseOneArg("executionMethod")
if err != nil {
logging.LogError(err, "failed to get executionMethod")
response.Success = false
response.Error = err.Error()
return response
Expand All @@ -160,6 +180,9 @@ func init() {
if remove {
displayParams += " -remove"
}
if executionMethod != "" {
displayParams += fmt.Sprintf(" -executionMethod %s", executionMethod)
}
response.DisplayParams = &displayParams
collectionSourceData, err := getCollectionSource(collection)
if err != nil {
Expand Down Expand Up @@ -207,6 +230,10 @@ func init() {
TaskID: taskData.Task.ID,
Response: []byte(fmt.Sprintf("Registering new command %s\n", prefixedCommandName)),
})
// Override execution method if provided via parameter
if executionMethod != "" {
commandSource.ExecutionMethod = executionMethod
}
newCommand := createAssemblyCommand(commandSource, collectionSourceData, true)
agentstructs.AllPayloadData.Get(PayloadTypeName).AddCommand(newCommand)
}
Expand Down