From 87b6d7faa4829b1e1c7c8895ef33d2b84d00b11f Mon Sep 17 00:00:00 2001
From: Mathieu Deous
Date: Sat, 6 May 2023 15:33:56 +0200
Subject: Better exit codes

---
 Makefile |  2 +-
 main.go  | 24 ++++++++++++++++++------
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/Makefile b/Makefile
index d5b9661..bd1ed51 100644
--- a/Makefile
+++ b/Makefile
@@ -38,7 +38,7 @@ docker: ## Build docker image
 	docker tag $(IMAGE_NAME):latest $(IMAGE_NAME):$(IMAGE_VERSION)
 
 docker-tests: ## Run docker image against the samples folder
-	@docker run --rm -v $(shell pwd)/data/samples:/data $(IMAGE_NAME):latest
+	@(docker run --rm -v $(shell pwd)/data/samples:/data $(IMAGE_NAME):latest && exit 1) || (test $$? -eq 255 || exit 1)
 
 docker-publish: ## Push docker image to the container registry
 	@docker push $(IMAGE_NAME):latest
diff --git a/main.go b/main.go
index de64b90..d4976eb 100644
--- a/main.go
+++ b/main.go
@@ -30,6 +30,9 @@ const (
 	FileBufferSize       = 32 * 1024 // 32KB
 	YaraMaxThreads       = 32
 	TempDirPrefix        = "pmf-"
+	ExitCodeOk           = 0
+	ExitCodeWithMatches  = 255
+	ExitCodeWithError    = 1
 )
 
 var (
@@ -93,7 +96,7 @@ func handleError(err error, desc string, isFatal bool) {
 		}
 		log.Println("[ERROR]"+desc, err)
 		if isFatal {
-			os.Exit(1)
+			os.Exit(ExitCodeWithError)
 		}
 	}
 }
@@ -111,8 +114,9 @@ func writeRulesFiles(content fs.FS) string {
 	for _, rulesFile := range rulesFiles {
 		// read embedded content
 		f, err := content.Open(path.Join("data", rulesFile))
-		handleError(err, "unable to read embedded rule", true)
+		handleError(err, "unable to open embedded rule", true)
 		ruleData, err := io.ReadAll(f)
+		handleError(err, "unable to read rule content", true)
 
 		// write to temporary file
 		err = os.WriteFile(path.Join(rulesPath, rulesFile), ruleData, 0640)
@@ -297,13 +301,14 @@ func loadRulesFile(fileName string) (*yara.Rules, error) {
 
 func main() {
 	startTime := time.Now()
+	matchesFound := false
 	_, err := flags.Parse(&args)
 	if err != nil {
-		os.Exit(1)
+		os.Exit(ExitCodeWithError)
 	}
 	if args.ShowVersion {
 		println(version)
-		os.Exit(0)
+		os.Exit(ExitCodeOk)
 	}
 
 	// check rules path
@@ -317,7 +322,7 @@ func main() {
 	// update rules if required
 	if args.Update {
 		updateRules()
-		os.Exit(0)
+		os.Exit(ExitCodeOk)
 	}
 
 	// add custom excluded file extensions
@@ -428,6 +433,7 @@ func main() {
 			for target, count := range matchCount {
 				if count >= DangerousMinScore {
 					log.Println("[WARNING] dangerous file found:", target)
+					matchesFound = true
 				}
 			}
 		} else { // single file mode
@@ -437,6 +443,7 @@ func main() {
 			err := scanner.SetCallback(&matches).ScanFile(args.Positional.Target)
 			handleError(err, "unable to scan target", true)
 			for _, match := range matches {
+				matchesFound = true
 				log.Println("[WARNING] match found:", args.Positional.Target, match.Rule)
 				if args.Verbose {
 					for _, matchString := range match.Strings {
@@ -460,6 +467,11 @@ func main() {
 			log.Println("[DEBUG] deleting temporary folder:", args.RulesDir)
 		}
 		err := os.RemoveAll(args.RulesDir)
-		handleError(err, "unable to delete temporary folder", true)
+		handleError(err, fmt.Sprintf("unable to delete temporary folder '%s'", args.RulesDir), false)
 	}
+
+	if matchesFound {
+		os.Exit(ExitCodeWithMatches)
+	}
+	os.Exit(ExitCodeOk)
 }
-- 
cgit v1.3