Skip to content

Commit

Permalink
Better exit codes
Browse files Browse the repository at this point in the history
  • Loading branch information
mdeous authored May 6, 2023
1 parent 1b85a73 commit 87b6d7f
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 7 deletions.
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
24 changes: 18 additions & 6 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,9 @@ const (
FileBufferSize = 32 * 1024 // 32KB
YaraMaxThreads = 32
TempDirPrefix = "pmf-"
ExitCodeOk = 0
ExitCodeWithMatches = 255
ExitCodeWithError = 1
)

var (
Expand Down Expand Up @@ -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)
}
}
}
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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 {
Expand All @@ -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)
}

0 comments on commit 87b6d7f

Please sign in to comment.