Fixed template expansion bug

Velocidex · Nov 27, 2023 · 4bdaf09 · 4bdaf09
1 parent f23468b
commit 4bdaf09
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 25 deletions.
diff --git a/Makefile b/Makefile
@@ -9,6 +9,9 @@ windows:
 	GOOS=windows GOARCH=amd64 go build -o velosigmac.exe ./src/*.go
 
 compile:
-	./velosigmac compile --config ./config/windows_hayabusa_rules.yaml --output ./output/Velociraptor-Hayabusa-Rules.zip
-	./velosigmac compile --config ./config/windows_hayabusa_event_monitoring.yaml --output ./output/Velociraptor-Hayabusa-Monitoring.zip
-	./velosigmac compile --config ./config/ChopChopGo_rules.yaml --output ./output/Velociraptor-ChopChopGo-Rules.zip
+	./velosigmac compile --config ./config/windows_hayabusa_rules.yaml --output ./output/Velociraptor-Hayabusa-Rules.zip --yaml ./output/Velociraptor-Hayabusa-Rules.yaml
+	./velosigmac compile --config ./config/windows_hayabusa_event_monitoring.yaml --output ./output/Velociraptor-Hayabusa-Monitoring.zip --yaml ./output/Velociraptor-Hayabusa-Monitoring.yaml
+	./velosigmac compile --config ./config/ChopChopGo_rules.yaml --output ./output/Velociraptor-ChopChopGo-Rules.zip --yaml ./output/Velociraptor-ChopChopGo-Rules.yaml
+
+test: compile
+	VELOCIRAPTOR_CONFIG= ../velociraptor/output/velociraptor-v0.7.1-rc1-linux-amd64 --definitions output/ artifacts list -v |grep 'Haya\|Chop'
diff --git a/config/windows_hayabusa_event_monitoring.yaml b/config/windows_hayabusa_event_monitoring.yaml
@@ -716,7 +716,7 @@ QueryTemplate: |
     FROM sigma(
       rules=split(string= Rules, sep_string="\n---\n"),
       log_sources= LogSources, debug=Debug,
-      default_details="{{.Base64DefaultDetailsQuery}}",
+      default_details='''{{.Base64DefaultDetailsQuery}}''',
       rule_filter="x=>x.Level =~ RuleLevelRegex AND x.Status =~ RuleStatusRegex AND x.Title =~ RuleTitleFilter",
       field_mapping= FieldMapping)
 

diff --git a/config/windows_hayabusa_rules.yaml b/config/windows_hayabusa_rules.yaml
@@ -716,7 +716,7 @@ QueryTemplate: |
     FROM sigma(
       rules=split(string= Rules, sep_string="\n---\n"),
       log_sources= LogSources, debug=Debug,
-      default_details="{{.Base64DefaultDetailsQuery}}",
+      default_details='''{{.Base64DefaultDetailsQuery}}''',
       rule_filter="x=>x.Level =~ RuleLevelRegex AND x.Status =~ RuleStatusRegex AND x.Title =~ RuleTitleFilter",
       field_mapping= FieldMapping)
 

diff --git a/src/artifact.go b/src/artifact.go
@@ -17,7 +17,7 @@ func encode(in []byte) string {
 	return base64.StdEncoding.EncodeToString(b.Bytes())
 }
 
-func (self *CompilerContext) WriteArtifact(zip *zip.Writer) error {
+func (self *CompilerContext) GetArtifact() (string, error) {
 	vql := BuildLogSource(self.config_obj)
 
 	params := &ArtifactContent{
@@ -29,22 +29,30 @@ func (self *CompilerContext) WriteArtifact(zip *zip.Writer) error {
 
 	templ, err := template.New("").Parse(self.config_obj.QueryTemplate)
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	b := &bytes.Buffer{}
 	err = templ.Execute(b, params)
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	vql += string(b.Bytes())
 
+	return self.config_obj.Preamble + indent(vql, 4), nil
+}
+
+func (self *CompilerContext) WriteArtifact(zip *zip.Writer) error {
+	artifact_yaml, err := self.GetArtifact()
+	if err != nil {
+		return err
+	}
 	fd, err := zip.Create("artifact.yaml")
 	if err != nil {
 		return err
 	}
-	fd.Write([]byte(self.config_obj.Preamble + indent(vql, 4)))
+	fd.Write([]byte(artifact_yaml))
 
 	// Also include the redacted rules in the zip file
 	fd, err = zip.Create("sigma_rules.yml")

diff --git a/src/compile.go b/src/compile.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"archive/zip"
+	"errors"
 	"fmt"
 	"io/fs"
 	"io/ioutil"
@@ -20,7 +21,8 @@ var (
 		"A tool for manipulating sigma files.")
 
 	compile_cmd     = app.Command("compile", "Compile all the rules into one rule.")
-	output          = compile_cmd.Flag("output", "File to write the artifact to").Required().String()
+	output          = compile_cmd.Flag("output", "File to write the artifact bundle to").String()
+	yaml_output     = compile_cmd.Flag("yaml", "File to write the artifact yaml to").String()
 	config          = compile_cmd.Flag("config", "Config file to use").Required().ExistingFile()
 	level_regex_str = compile_cmd.Flag("level_regex", "A regex to select rule Levels").Default(".").String()
 
@@ -118,24 +120,23 @@ func (self *CompilerContext) CompileDirs() error {
 	return nil
 }
 
-func doCompile() {
-	// Write the sigma file in the output directory.
-	out_fd, err := os.OpenFile(*output,
-		os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
-	kingpin.FatalIfError(err, "Creating output")
-	defer out_fd.Close()
+func doCompile() error {
+	if *yaml_output == "" && *output == "" {
+		return errors.New("Must provide either --output or --yaml")
+	}
 
 	level_regex, err := regexp.Compile(*level_regex_str)
-	kingpin.FatalIfError(err, "Level Regex invalid")
-
-	zip := zip.NewWriter(out_fd)
-	defer zip.Close()
+	if err != nil {
+		return fmt.Errorf("Level Regex invalid: %w", err)
+	}
 
 	context := NewCompilerContext()
 	context.level_regex = level_regex
 
 	err = context.LoadConfig(*config)
-	kingpin.FatalIfError(err, "Reading Config")
+	if err != nil {
+		return fmt.Errorf("Reading Config: %w", err)
+	}
 
 	defer func() {
 		fmt.Printf("Generated rules with level %v into %v\n",
@@ -144,17 +145,52 @@ func doCompile() {
 	}()
 
 	err = context.CompileDirs()
-	kingpin.FatalIfError(err, "Listing directory")
+	if err != nil {
+		return fmt.Errorf("Listing directory: %w", err)
+	}
 
-	err = context.WriteArtifact(zip)
-	kingpin.FatalIfError(err, "WriteArtifact")
+	if *output != "" {
+		// Write the sigma file in the output directory.
+		out_fd, err := os.OpenFile(*output,
+			os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
+		if err != nil {
+			return fmt.Errorf("Creating output: %w", err)
+		}
+		defer out_fd.Close()
+
+		zip := zip.NewWriter(out_fd)
+		defer zip.Close()
+
+		err = context.WriteArtifact(zip)
+		if err != nil {
+			return fmt.Errorf("WriteArtifact: %w", err)
+		}
+	}
+
+	if *yaml_output != "" {
+		out_fd, err := os.OpenFile(*yaml_output,
+			os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
+		if err != nil {
+			return fmt.Errorf("Creating yaml output: %w", err)
+		}
+		defer out_fd.Close()
+
+		artifact, err := context.GetArtifact()
+		if err != nil {
+			return fmt.Errorf("GetArtifact: %w", err)
+		}
+
+		out_fd.Write([]byte(artifact))
+	}
+	return nil
 }
 
 func init() {
 	command_handlers = append(command_handlers, func(command string) bool {
 		switch command {
 		case compile_cmd.FullCommand():
-			doCompile()
+			err := doCompile()
+			kingpin.FatalIfError(err, "Compiling artifact")
 
 		default:
 			return false