Feature/update variables api (#35)

Co-authored-by: Hidde-Jan Jongsma <hidde-jan.jongsma@tno.nl>

internal/capability/capability.go

@@ -10,6 +10,6 @@ type ICapability interface {
 		command cacao.Command,
 		authentication cacao.AuthenticationInformation,
 		target cacao.AgentTarget,
-		variables cacao.VariableMap) (cacao.VariableMap, error)
+		variables cacao.Variables) (cacao.Variables, error)
 	GetType() string
 }

internal/capability/http/http.go

@@ -38,13 +38,13 @@ func (httpCapability *HttpCapability) Execute(
 	command cacao.Command,
 	authentication cacao.AuthenticationInformation,
 	target cacao.AgentTarget,
-	variables cacao.VariableMap) (cacao.VariableMap, error) {
+	variables cacao.Variables) (cacao.Variables, error) {
 
 	// Get request data and handle errors
 	method, url, errmethod := ObtainHttpMethodAndUrlFromCommand(command)
 	if errmethod != nil {
 		log.Error(errmethod)
-		return cacao.VariableMap{}, errmethod
+		return cacao.Variables{}, errmethod
 	}
 	content_data, errcontent := ObtainHttpRequestContentDataFromCommand(command)
 	if errcontent != nil {
@@ -56,7 +56,7 @@ func (httpCapability *HttpCapability) Execute(
 	request, err := http.NewRequest(method, url, bytes.NewBuffer(content_data))
 	if err != nil {
 		log.Error(err)
-		return cacao.VariableMap{}, err
+		return cacao.Variables{}, err
 	}
 
 	for key, httpCapability := range command.Headers {
@@ -65,12 +65,12 @@ func (httpCapability *HttpCapability) Execute(
 	if target.ID != "" {
 		if err := verifyAuthInfoMatchesAgentTarget(&target, &authentication); err != nil {
 			log.Error(err)
-			return cacao.VariableMap{}, err
+			return cacao.Variables{}, err
 		}
 
 		if err := setupAuthHeaders(request, &authentication); err != nil {
 			log.Error(err)
-			return cacao.VariableMap{}, err
+			return cacao.Variables{}, err
 		}
 	}
 
@@ -79,22 +79,22 @@ func (httpCapability *HttpCapability) Execute(
 	response, err := client.Do(request)
 	if err != nil {
 		log.Error(err)
-		return cacao.VariableMap{}, err
+		return cacao.Variables{}, err
 	}
 	defer response.Body.Close()
 
 	responseBytes, err := io.ReadAll(response.Body)
 	if err != nil {
 		log.Error(err)
-		return cacao.VariableMap{}, err
+		return cacao.Variables{}, err
 	}
 	respString := string(responseBytes)
 	sc := response.StatusCode
 	if sc < 200 || sc > 299 {
-		return cacao.VariableMap{}, errors.New(respString)
+		return cacao.Variables{}, errors.New(respString)
 	}
 
-	return cacao.VariableMap{
+	return cacao.Variables{
 		"__soarca_http_result__": {Name: "result", Value: respString}}, nil
 
 }

internal/capability/ssh/ssh.go

@@ -31,7 +31,7 @@ func (sshCapability *SshCapability) Execute(metadata execution.Metadata,
 	command cacao.Command,
 	authentication cacao.AuthenticationInformation,
 	target cacao.AgentTarget,
-	variables cacao.VariableMap) (cacao.VariableMap, error) {
+	variables cacao.Variables) (cacao.Variables, error) {
 	log.Trace(metadata.ExecutionId)
 
 	host := CombinePortAndAddress(target.Address, target.Port)
@@ -40,7 +40,7 @@ func (sshCapability *SshCapability) Execute(metadata execution.Metadata,
 
 	if errAuth != nil {
 		log.Error(errAuth)
-		return cacao.VariableMap{}, errAuth
+		return cacao.Variables{}, errAuth
 	} else {
 		log.Trace(host)
 	}
@@ -60,7 +60,7 @@ func (sshCapability *SshCapability) Execute(metadata execution.Metadata,
 		signer, errKey := ssh.ParsePrivateKey([]byte(authentication.PrivateKey))
 		if errKey != nil || authentication.Password == "" {
 			log.Error("no valid authentication information: ", errKey)
-			return cacao.VariableMap{}, errKey
+			return cacao.Variables{}, errKey
 		}
 		config = ssh.ClientConfig{
 			User: authentication.Username,
@@ -78,23 +78,23 @@ func (sshCapability *SshCapability) Execute(metadata execution.Metadata,
 	conn, err := ssh.Dial("tcp", host, &config)
 	if err != nil {
 		log.Error(err)
-		return cacao.VariableMap{}, err
+		return cacao.Variables{}, err
 	}
 	var session *ssh.Session
 	session, err = conn.NewSession()
 	if err != nil {
 		log.Error(err)
-		return cacao.VariableMap{}, err
+		return cacao.Variables{}, err
 	}
 
 	response, err := session.Output(StripSshPrepend(command.Command))
 	defer session.Close()
 
 	if err != nil {
 		log.Error(err)
-		return cacao.VariableMap{"__soarca_ssh_result__": {Name: "result", Value: string(response)}}, err
+		return cacao.Variables{"__soarca_ssh_result__": {Name: "result", Value: string(response)}}, err
 	}
-	results := cacao.VariableMap{"__soarca_ssh_result__": {Name: "result", Value: string(response)}}
+	results := cacao.Variables{"__soarca_ssh_result__": {Name: "result", Value: string(response)}}
 	log.Trace("Finished ssh execution will return the variables: ", results)
 	return results, err
 }

internal/decomposer/decomposer.go

@@ -49,7 +49,7 @@ type Decomposer struct {
 	guid      guid.IGuid
 }
 
-func Callback(executionId uuid.UUID, outputVariables cacao.VariableMap) {
+func Callback(executionId uuid.UUID, outputVariables cacao.Variables) {
 
 }
 

internal/executer/executer.go

@@ -21,8 +21,8 @@ type IExecuter interface {
 		command cacao.Command,
 		authentication cacao.AuthenticationInformation,
 		target cacao.AgentTarget,
-		variable cacao.VariableMap,
-		module cacao.AgentTarget) (uuid.UUID, cacao.VariableMap, error)
+		variable cacao.Variables,
+		module cacao.AgentTarget) (uuid.UUID, cacao.Variables, error)
 }
 
 func init() {
@@ -43,14 +43,14 @@ func (executer *Executer) Execute(metadata execution.Metadata,
 	command cacao.Command,
 	authentication cacao.AuthenticationInformation,
 	target cacao.AgentTarget,
-	variable cacao.VariableMap,
-	agent cacao.AgentTarget) (uuid.UUID, cacao.VariableMap, error) {
+	variable cacao.Variables,
+	agent cacao.AgentTarget) (uuid.UUID, cacao.Variables, error) {
 
 	if capability, ok := executer.capabilities[agent.Name]; ok {
 		returnVariables, err := capability.Execute(metadata, command, authentication, target, variable)
 		return metadata.ExecutionId, returnVariables, err
 	} else {
-		empty := cacao.VariableMap{}
+		empty := cacao.Variables{}
 		message := "executor is not available in soarca"
 		err := errors.New(message)
 		log.Error(message)

models/cacao/cacao.go

@@ -29,6 +29,17 @@ type (
 	Workflow   map[string]Step
 )
 
+type Variable struct {
+	Type        string `bson:"type" json:"type" validate:"required"`
+	Name        string `bson:"name,omitempty" json:"name,omitempty"`
+	Description string `bson:"description,omitempty" json:"description,omitempty"`
+	Value       string `bson:"value,omitempty" json:"value,omitempty"`
+	Constant    bool   `bson:"constant,omitempty" json:"constant,omitempty"`
+	External    bool   `bson:"external,omitempty" json:"external,omitempty"`
+}
+
+type Variables map[string]Variable
+
 type Playbook struct {
 	ID                            string                               `bson:"_id" json:"id" validate:"required"`
 	Type                          string                               `bson:"type" json:"type" validate:"required"`
@@ -56,7 +67,7 @@ type Playbook struct {
 	AgentDefinitions              map[string]AgentTarget               `bson:"agent_definitions,omitempty" json:"agent_definitions,omitempty"`
 	TargetDefinitions             map[string]AgentTarget               `bson:"target_definitions,omitempty" json:"target_definitions,omitempty"`
 	ExtensionDefinitions          map[string]ExtensionDefinition       `bson:"extension_definitions,omitempty" json:"extension_definitions,omitempty"`
-	PlaybookVariables             VariableMap                          `bson:"playbook_variables,omitempty" json:"playbook_variables,omitempty"`
+	PlaybookVariables             Variables                            `bson:"playbook_variables,omitempty" json:"playbook_variables,omitempty"`
 	PlaybookExtensions            Extensions                           `bson:"playbook_extensions,omitempty" json:"playbook_extensions,omitempty"`
 }
 
@@ -151,7 +162,7 @@ type Step struct {
 	ExternalReferences []ExternalReferences `bson:"external_references,omitempty" json:"external_references,omitempty"`
 	Delay              int                  `bson:"delay,omitempty" json:"delay,omitempty"`
 	Timeout            int                  `bson:"timeout,omitempty" json:"timeout,omitempty"`
-	StepVariables      VariableMap          `bson:"step_variables,omitempty" json:"step_variables,omitempty"`
+	StepVariables      Variables            `bson:"step_variables,omitempty" json:"step_variables,omitempty"`
 	Owner              string               `bson:"owner,omitempty" json:"owner,omitempty"`
 	OnCompletion       string               `bson:"on_completion,omitempty" json:"on_completion,omitempty"`
 	OnSuccess          string               `bson:"on_success,omitempty" json:"on_success,omitempty"`

models/cacao/variables.go

@@ -5,39 +5,82 @@ import (
 	"strings"
 )
 
-type Variable struct {
-	Type        string `bson:"type" json:"type" validate:"required"`
-	Name        string `bson:"name,omitempty" json:"name,omitempty"`
-	Description string `bson:"description,omitempty" json:"description,omitempty"`
-	Value       string `bson:"value,omitempty" json:"value,omitempty"`
-	Constant    bool   `bson:"constant,omitempty" json:"constant,omitempty"`
-	External    bool   `bson:"external,omitempty" json:"external,omitempty"`
+// Initialize new Variables
+//
+// Allows passing in multiple variables at once
+func NewVariables(new ...Variable) Variables {
+	variables := make(Variables)
+	for _, variable := range new {
+		variables.Insert(variable)
+	}
+	return variables
+}
+
+// Insert a variable
+//
+// Returns true if the variable was inserted
+func (variables *Variables) Insert(new Variable) bool {
+	if _, found := (*variables)[new.Name]; found {
+		return false
+	}
+	(*variables)[new.Name] = new
+	return true
+}
+
+// Insert or replace a variable
+//
+// Returns true if the variable was replaced
+func (variables *Variables) InsertOrReplace(new Variable) bool {
+	_, found := (*variables)[new.Name]
+	(*variables)[new.Name] = new
+	return found
+}
+
+// Insert variables map at once into the base and keep base variables in favor of source duplicates
+func (variables *Variables) InsertRange(source Variables) {
+	for _, variable := range source {
+		variables.Insert(variable)
+	}
 }
 
-type VariableMap map[string]Variable
+// Find a variable by name
+//
+// Returns a Variable struct and a boolean indicating if it was found
+func (variables Variables) Find(key string) (Variable, bool) {
+	val, ok := variables[key]
+	return val, ok
+}
 
-func (variableMap VariableMap) Merge(otherMap VariableMap) VariableMap {
-	newMap := make(VariableMap)
-	for k, v := range variableMap {
-		newMap[k] = v
+// Interpolate variable references into a target string
+//
+// Returns the Interpolated string with variables values available in the map
+func (variables *Variables) Interpolate(input string) string {
+	replacements := make([]string, 0)
+	for key, value := range *variables {
+		replacementKey := fmt.Sprint(key, ":value")
+		replacements = append(replacements, replacementKey, value.Value)
 	}
+	return strings.NewReplacer(replacements...).Replace(input)
+}
+
+// Select a subset of variables from the map
+//
+// Unknown keys are ignored.
+func (variables *Variables) Select(keys []string) Variables {
+	newVariables := NewVariables()
 
-	for k, v := range otherMap {
-		if _, ok := variableMap[k]; !ok || v.Constant {
-			newMap[k] = v
+	for _, key := range keys {
+		if value, ok := variables.Find(key); ok {
+			newVariables.InsertOrReplace(value)
 		}
 	}
 
-	return newMap
+	return newVariables
 }
 
-// TODO: Find a way to store the Replacer 'inside' the VariableMap
-func (variableMap VariableMap) Replace(s string) string {
-	replacements := make([]string, 0)
-	for k, v := range variableMap {
-		replacementKey := fmt.Sprintf("%s:value", k)
-		replacements = append(replacements, replacementKey, v.Value)
+// Merge two maps of Cacao variables and replace the base with the source if exists
+func (variables *Variables) Merge(source Variables) {
+	for _, variable := range source {
+		variables.InsertOrReplace(variable)
 	}
-	r := strings.NewReplacer(replacements...)
-	return r.Replace(s)
 }

models/decoder/cacao.go

@@ -68,5 +68,17 @@ func decode(data []byte) *cacao.Playbook {
 		playbook.AuthenticationInfoDefinitions[key] = auth
 	}
 
+	for key, variable := range playbook.PlaybookVariables {
+		variable.Name = key
+		playbook.PlaybookVariables.InsertOrReplace(variable)
+	}
+
+	for _, step := range playbook.Workflow {
+		for key, variable := range step.StepVariables {
+			variable.Name = key
+			step.StepVariables.InsertOrReplace(variable)
+		}
+	}
+
 	return &playbook
 }

test/integration/capability/http/http_integration_test.go

@@ -35,7 +35,7 @@ func TestHttpConnection(t *testing.T) {
 		metadata, expectedCommand,
 		cacao.AuthenticationInformation{},
 		cacao.AgentTarget{},
-		cacao.VariableMap{"test": variable1})
+		cacao.Variables{"test": variable1})
 	if err != nil {
 		fmt.Println(err)
 		t.Fail()
@@ -81,7 +81,7 @@ func TestHttpOAuth2(t *testing.T) {
 		expectedCommand,
 		oauth2_info,
 		target,
-		cacao.VariableMap{"test": variable1})
+		cacao.Variables{"test": variable1})
 	if err != nil {
 		fmt.Println(err)
 		t.Fail()
@@ -125,7 +125,7 @@ func TestHttpBasicAuth(t *testing.T) {
 		metadata,
 		expectedCommand,
 		basicauth_info,
-		target, cacao.VariableMap{"test": variable1})
+		target, cacao.Variables{"test": variable1})
 	if err != nil {
 		fmt.Println(err)
 		t.Fail()

test/integration/capability/ssh/ssh_integration_test.go

@@ -45,7 +45,7 @@ func TestSshConnection(t *testing.T) {
 		expectedCommand,
 		expectedAuthenticationInformation,
 		expectedTarget,
-		cacao.VariableMap{expectedVariables.Name: expectedVariables})
+		cacao.Variables{expectedVariables.Name: expectedVariables})
 	if err != nil {
 		fmt.Println(err)
 		t.Fail()