working on integrating script engine into threagile

Author: joreiche

cmd/risk_demo/main.go

@@ -49,7 +49,12 @@ func main() {
 			os.Exit(-2)
 		}
 
-		generatedRisks := new(customRiskRule).GenerateRisks(&input)
+		generatedRisks, riskError := new(customRiskRule).GenerateRisks(&input)
+		if riskError != nil {
+			_, _ = fmt.Fprintf(os.Stderr, "failed to generate risks: %v\n", riskError)
+			os.Exit(-2)
+		}
+
 		outData, marshalError := json.Marshal(generatedRisks)
 		if marshalError != nil {
 			_, _ = fmt.Fprintf(os.Stderr, "failed to print generated risks: %v\n", marshalError)
@@ -89,17 +94,17 @@ func (r customRiskRule) SupportedTags() []string {
 	return []string{"demo tag"}
 }
 
-func (r customRiskRule) GenerateRisks(parsedModel *types.Model) []types.Risk {
-	generatedRisks := make([]types.Risk, 0)
+func (r customRiskRule) GenerateRisks(parsedModel *types.Model) ([]*types.Risk, error) {
+	generatedRisks := make([]*types.Risk, 0)
 	for _, techAsset := range parsedModel.TechnicalAssets {
 		generatedRisks = append(generatedRisks, createRisk(techAsset))
 	}
-	return generatedRisks
+	return generatedRisks, nil
 }
 
-func createRisk(technicalAsset *types.TechnicalAsset) types.Risk {
+func createRisk(technicalAsset *types.TechnicalAsset) *types.Risk {
 	category := new(customRiskRule).Category()
-	risk := types.Risk{
+	risk := &types.Risk{
 		CategoryId:                   category.ID,
 		Severity:                     types.CalculateSeverity(types.VeryLikely, types.MediumImpact),
 		ExploitationLikelihood:       types.VeryLikely,

cmd/script/main.go

@@ -2,99 +2,52 @@ package main
 
 import (
 	"fmt"
-	"github.com/threagile/threagile/pkg/common"
-	"github.com/threagile/threagile/pkg/input"
-	"github.com/threagile/threagile/pkg/model"
 	"github.com/threagile/threagile/pkg/script"
-	"github.com/threagile/threagile/pkg/security/risks"
+	"github.com/threagile/threagile/pkg/security/types"
 	"gopkg.in/yaml.v3"
 	"os"
 	"path/filepath"
 )
 
 func main() {
-	ruleData, ruleReadError := os.ReadFile(filepath.Join("test", "risk-category.yaml"))
+	scriptFilename := filepath.Join("test", "risk-category.yaml")
+	ruleData, ruleReadError := os.ReadFile(scriptFilename)
 	if ruleReadError != nil {
 		fmt.Printf("error reading risk category: %v\n", ruleReadError)
 		return
 	}
 
-	scripts, parseError := new(script.Script).ParseScripts(ruleData)
+	newRule, parseError := new(script.RiskRule).ParseFromData(ruleData)
 	if parseError != nil {
-		fmt.Printf("error parsing scripts: %v\n", parseError)
+		fmt.Printf("error parsing scripts from %q: %v\n", scriptFilename, parseError)
 		return
 	}
 
-	modelData, modelReadError := os.ReadFile(filepath.Join("test", "parsed-model.yaml"))
+	modelFilename := filepath.Join("test", "parsed-model.yaml")
+	modelData, modelReadError := os.ReadFile(modelFilename)
 	if modelReadError != nil {
 		fmt.Printf("error reading model: %v\n", modelReadError)
 		return
 	}
 
-	inputModel := new(input.Model)
-	modelUnmarshalError := yaml.Unmarshal(modelData, inputModel)
+	parsedModel := new(types.Model)
+	modelUnmarshalError := yaml.Unmarshal(modelData, parsedModel)
 	if modelUnmarshalError != nil {
-		fmt.Printf("error parsing model: %v\n", modelUnmarshalError)
+		fmt.Printf("error parsing model from %q: %v\n", modelFilename, modelUnmarshalError)
 		return
 	}
 
-	/*
-		categoriesModel := new(input.Model)
-		riskUnmarshalError := yaml.Unmarshal(riskData, categoriesModel)
-		if riskUnmarshalError != nil {
-			fmt.Printf("error parsing risk category: %v\n", riskUnmarshalError)
-			return
-		}
-	*/
-
-	parsedModel, modelError := model.ParseModel(&common.Config{}, inputModel, make(risks.RiskRules), make(risks.RiskRules))
-	if modelError != nil {
-		fmt.Printf("error importing model: %v\n", modelError)
+	risks, riskError := newRule.GenerateRisks(parsedModel)
+	if riskError != nil {
+		fmt.Printf("error generating risks for %q: %v\n", newRule.Category().ID, riskError)
 		return
 	}
 
-	_ = parsedModel
-	_ = scripts
-	/*
-		var risk types.RiskCategory
-		if categoriesModel.CustomRiskCategories != nil {
-			for _, item := range categoriesModel.CustomRiskCategories {
-				risk = item
-			}
-		}
-
-		if len(categoriesModel.CustomRiskCategories) == 0 {
-			fmt.Printf("no risk categories\n")
-			return
-		}
-
-		for name, script := range scripts {
-			scope := new(script.Scope)
-			addError := scope.Init(parsedModel, &risk, script.Utils())
-			if addError != nil {
-				fmt.Printf("error adding model to scope for %q: %v\n", name, addError)
-				return
-			}
-
-			risks, errorLiteral, riskError := script.GenerateRisks(scope)
-			if riskError != nil {
-				fmt.Printf("error generating risks for %q: %v\n", name, riskError)
-
-				if len(errorLiteral) > 0 {
-					fmt.Printf("in:\n%v\n", script.IndentPrintf(1, errorLiteral))
-				}
-
-				return
-			}
-
-			printedRisks, printError := yaml.Marshal(risks)
-			if printError != nil {
-				fmt.Printf("error printing risks for %q: %v\n", name, printError)
-				return
-			}
-
-			fmt.Printf("generated risks for %q: \n%v\n", name, string(printedRisks))
-		}
+	printedRisks, printError := yaml.Marshal(risks)
+	if printError != nil {
+		fmt.Printf("error printing risks for %q: %v\n", newRule.Category().ID, printError)
+		return
+	}
 
-	*/
+	fmt.Printf("generated risks for %q: \n%v\n", newRule.Category().ID, string(printedRisks))
 }

pkg/common/config.go

@@ -2,6 +2,7 @@ package common
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -157,39 +158,49 @@ func (c *Config) Load(configFilename string) error {
 
 	c.Merge(config, values)
 
+	errorList := make([]error, 0)
 	c.TempFolder = c.CleanPath(c.TempFolder)
 	tempDirError := os.MkdirAll(c.TempFolder, 0700)
 	if tempDirError != nil {
-		return fmt.Errorf("failed to create temp dir %q: %v", c.TempFolder, tempDirError)
+		errorList = append(errorList, fmt.Errorf("failed to create temp dir %q: %v", c.TempFolder, tempDirError))
 	}
 
 	c.OutputFolder = c.CleanPath(c.OutputFolder)
 	outDirError := os.MkdirAll(c.OutputFolder, 0700)
 	if outDirError != nil {
-		return fmt.Errorf("failed to create output dir %q: %v", c.OutputFolder, outDirError)
+		errorList = append(errorList, fmt.Errorf("failed to create output dir %q: %v", c.OutputFolder, outDirError))
 	}
 
 	c.AppFolder = c.CleanPath(c.AppFolder)
 	appDirError := c.checkDir(c.AppFolder, "app")
 	if appDirError != nil {
-		return appDirError
+		errorList = append(errorList, appDirError)
 	}
 
 	c.PluginFolder = c.CleanPath(c.PluginFolder)
 	binDirError := c.checkDir(c.PluginFolder, "plugin")
 	if binDirError != nil {
-		return binDirError
+		errorList = append(errorList, binDirError)
 	}
 
 	c.DataFolder = c.CleanPath(c.DataFolder)
 	dataDirError := c.checkDir(c.DataFolder, "data")
 	if dataDirError != nil {
-		return dataDirError
+		errorList = append(errorList, dataDirError)
 	}
 
 	c.TechnologyFilename = c.CleanPath(c.TechnologyFilename)
 
-	return c.CheckServerFolder()
+	serverFolderError := c.CheckServerFolder()
+	if serverFolderError != nil {
+		errorList = append(errorList, serverFolderError)
+	}
+
+	if len(errorList) > 0 {
+		return errors.Join(errorList...)
+	}
+
+	return nil
 }
 
 func (c *Config) CheckServerFolder() error {

pkg/model/custom-risk-category.go

@@ -3,7 +3,6 @@ package model
 import (
 	"fmt"
 	"github.com/threagile/threagile/pkg/security/risks"
-	"log"
 	"strings"
 
 	"github.com/threagile/threagile/pkg/security/types"
@@ -33,18 +32,18 @@ func (what *CustomRiskCategory) SupportedTags() []string {
 	return what.Tags
 }
 
-func (what *CustomRiskCategory) GenerateRisks(parsedModel *types.Model) []*types.Risk {
+func (what *CustomRiskCategory) GenerateRisks(parsedModel *types.Model) ([]*types.Risk, error) {
 	if what.runner == nil {
-		return nil
+		return nil, nil
 	}
 
 	generatedRisks := make([]*types.Risk, 0)
 	runError := what.runner.Run(parsedModel, &generatedRisks, "-generate-risks")
 	if runError != nil {
-		log.Fatalf("Failed to generate risks for custom risk rule %q: %v\n", what.runner.Filename, runError)
+		return nil, fmt.Errorf("Failed to generate risks for custom risk rule %q: %v\n", what.runner.Filename, runError)
 	}
 
-	return generatedRisks
+	return generatedRisks, nil
 }
 
 func LoadCustomRiskRules(pluginFiles []string, reporter types.ProgressReporter) risks.RiskRules {

pkg/model/read.go

@@ -43,6 +43,14 @@ func ReadAndAnalyzeModel(config *common.Config, progressReporter types.ProgressR
 		return nil, fmt.Errorf("unable to parse model yaml: %v", parseError)
 	}
 
+	/**
+	jsonData, _ := json.MarshalIndent(parsedModel, "", "  ")
+	_ = os.WriteFile("parsed-model.json", jsonData, 0600)
+
+	yamlData, _ := yaml.Marshal(parsedModel)
+	_ = os.WriteFile("parsed-model.yaml", yamlData, 0600)
+	/**/
+
 	introTextRAA := applyRAA(parsedModel, config.PluginFolder, config.RAAPlugin, progressReporter)
 
 	applyRiskGeneration(parsedModel, builtinRiskRules.Merge(customRiskRules), config.SkipRiskRules, progressReporter)
@@ -56,14 +64,6 @@ func ReadAndAnalyzeModel(config *common.Config, progressReporter types.ProgressR
 		return nil, fmt.Errorf("unable to check risk tracking: %v", err)
 	}
 
-	/*
-		jsonData, _ := json.MarshalIndent(parsedModel, "", "  ")
-		_ = os.WriteFile("parsed-model.json", jsonData, 0600)
-
-		yamlData, _ := yaml.Marshal(parsedModel)
-		_ = os.WriteFile("parsed-model.yaml", yamlData, 0600)
-	*/
-
 	return &ReadResult{
 		ModelInput:       modelInput,
 		ParsedModel:      parsedModel,
@@ -94,7 +94,12 @@ func applyRiskGeneration(parsedModel *types.Model, rules risks.RiskRules,
 		}
 
 		parsedModel.AddToListOfSupportedTags(rule.SupportedTags())
-		newRisks := rule.GenerateRisks(parsedModel)
+		newRisks, riskError := rule.GenerateRisks(parsedModel)
+		if riskError != nil {
+			progressReporter.Warnf("Error generating risks for %q: %v", id, riskError)
+			continue
+		}
+
 		if len(newRisks) > 0 {
 			parsedModel.GeneratedRisksByCategory[id] = newRisks
 		}

pkg/script/common/key-words.go

@@ -1,9 +1,10 @@
 package common
 
 const (
-	Risk  = "risk"
-	Match = "match"
-	Utils = "utils"
+	Script = "script"
+	Risk   = "risk"
+	Match  = "match"
+	Utils  = "utils"
 
 	Assign = "assign"
 	Loop   = "loop"

pkg/script/common/scope.go

@@ -16,33 +16,37 @@ type Scope struct {
 	returnValue Value
 }
 
-func (what *Scope) Init(model *types.Model, risk *types.RiskCategory, methods map[string]Statement) error {
-	if model != nil {
-		data, marshalError := json.Marshal(model)
+func (what *Scope) Init(risk *types.RiskCategory, methods map[string]Statement) error {
+	if risk != nil {
+		data, marshalError := json.Marshal(risk)
 		if marshalError != nil {
 			return marshalError
 		}
 
-		unmarshalError := json.Unmarshal(data, &what.Model)
+		unmarshalError := json.Unmarshal(data, &what.Risk)
 		if unmarshalError != nil {
 			return unmarshalError
 		}
 	}
 
-	if risk != nil {
-		data, marshalError := json.Marshal(risk)
+	what.Methods = methods
+
+	return nil
+}
+
+func (what *Scope) SetModel(model *types.Model) error {
+	if model != nil {
+		data, marshalError := json.Marshal(model)
 		if marshalError != nil {
 			return marshalError
 		}
 
-		unmarshalError := json.Unmarshal(data, &what.Risk)
+		unmarshalError := json.Unmarshal(data, &what.Model)
 		if unmarshalError != nil {
 			return unmarshalError
 		}
 	}
 
-	what.Methods = methods
-
 	return nil
 }