Skip to content

Commit 49cc63a

Browse files
xfhgxfhg
committed
v1.0.X
1 parent 44822a3 commit 49cc63a

File tree

6 files changed

+291
-181
lines changed

6 files changed

+291
-181
lines changed

‎cmd/assure.go

+122-28
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"os"
77
"os/exec"
88
"path/filepath"
9+
"sync"
910
)
1011

1112
// ProcessAssureType handles the assurance process for policies of type "assure"
@@ -64,39 +65,19 @@ func executeAssure(policy Policy, rgPath string, targetDir string, filesToAssure
6465
return fmt.Errorf("no target directory defined")
6566
}
6667

67-
// Append the file targets
68-
if len(filesToAssure) > 0 {
69-
codePatternAssureJSON = append(codePatternAssureJSON, filesToAssure...)
70-
} else if policy.FilePattern == "" {
71-
codePatternAssureJSON = append(codePatternAssureJSON, targetDir)
68+
matchesFound := true
69+
70+
// Parallel execution for large file sets
71+
if len(filesToAssure) > 25 {
72+
matchesFound, err = executeParallelAssure(rgPath, codePatternAssureJSON, filesToAssure, writer)
7273
} else {
73-
return fmt.Errorf("no files matched policy pattern")
74+
matchesFound, err = executeSingleAssure(rgPath, codePatternAssureJSON, filesToAssure, targetDir, policy, writer)
7475
}
7576

76-
// Execute the ripgrep command for JSON output
77-
cmdJSON := exec.Command(rgPath, codePatternAssureJSON...)
78-
cmdJSON.Stdout = writer
79-
cmdJSON.Stderr = os.Stderr
77+
if err != nil {
8078

81-
log.Debug().Msgf("Creating JSON output for assure policy %s... ", policy.ID)
82-
err = cmdJSON.Run()
79+
log.Error().Err(err).Msg("error executing ripgrep batch")
8380

84-
// Check if ripgrep found any matches
85-
matchesFound := true
86-
if err != nil {
87-
if exitError, ok := err.(*exec.ExitError); ok {
88-
// Exit code 1 in ripgrep means "no matches found"
89-
if exitError.ExitCode() == 1 {
90-
matchesFound = false
91-
err = nil // Reset error as this is the expected outcome for assure
92-
} else {
93-
log.Error().Err(err).Msg("error executing ripgrep for JSON output")
94-
return fmt.Errorf("error executing ripgrep for JSON output: %w", err)
95-
}
96-
} else {
97-
log.Error().Err(err).Msg("error executing ripgrep for JSON output")
98-
return fmt.Errorf("error executing ripgrep for JSON output: %w", err)
99-
}
10081
}
10182

10283
// Patch the JSON output file
@@ -107,6 +88,7 @@ func executeAssure(policy Policy, rgPath string, targetDir string, filesToAssure
10788
}
10889

10990
log.Debug().Msgf("JSON output for assure policy %s written to: %s ", policy.ID, jsonOutputFile)
91+
log.Debug().Msgf("Scanned ~%d files for policy %s", len(filesToAssure), policy.ID)
11092

11193
// Determine the status based on whether matches were found
11294
status := "NOT FOUND"
@@ -143,3 +125,115 @@ func executeAssure(policy Policy, rgPath string, targetDir string, filesToAssure
143125

144126
return nil
145127
}
128+
129+
func executeParallelAssure(rgPath string, baseArgs []string, filesToScan []string, writer *bufio.Writer) (bool, error) {
130+
131+
const batchSize = 25
132+
matched := true
133+
var wg sync.WaitGroup
134+
errChan := make(chan error, len(filesToScan)/batchSize+1)
135+
var mu sync.Mutex
136+
137+
for i := 0; i < len(filesToScan); i += batchSize {
138+
end := i + batchSize
139+
if end > len(filesToScan) {
140+
end = len(filesToScan)
141+
}
142+
batch := filesToScan[i:end]
143+
144+
// log.Debug().Msgf("RGM: %v", batch)
145+
146+
wg.Add(1)
147+
go func(batch []string) {
148+
defer wg.Done()
149+
args := append(baseArgs, batch...)
150+
cmd := exec.Command(rgPath, args...)
151+
output, err := cmd.Output()
152+
153+
if err != nil {
154+
155+
if exitError, ok := err.(*exec.ExitError); ok {
156+
// Exit code 1 in ripgrep means "no matches found"
157+
if exitError.ExitCode() == 1 {
158+
matched = false
159+
err = nil // Reset error as this is the expected outcome for assure
160+
}
161+
}
162+
163+
if exitError, ok := err.(*exec.ExitError); ok && exitError.ExitCode() != 1 {
164+
errChan <- fmt.Errorf("error executing ripgrep: %w", err)
165+
return
166+
}
167+
}
168+
169+
mu.Lock()
170+
_, writeErr := writer.Write(output)
171+
if writeErr == nil {
172+
writeErr = writer.Flush()
173+
}
174+
mu.Unlock()
175+
176+
if writeErr != nil {
177+
errChan <- fmt.Errorf("error writing output: %w", writeErr)
178+
}
179+
}(batch)
180+
}
181+
182+
wg.Wait()
183+
close(errChan)
184+
185+
for err := range errChan {
186+
if err != nil {
187+
return matched, err
188+
}
189+
}
190+
191+
return matched, nil
192+
}
193+
194+
func executeSingleAssure(rgPath string, baseArgs []string, filesToScan []string, targetDir string, policy Policy, writer *bufio.Writer) (bool, error) {
195+
196+
if len(filesToScan) > 0 {
197+
baseArgs = append(baseArgs, filesToScan...)
198+
} else {
199+
log.Error().Str("policy", policy.ID).Msgf("no files matched policy pattern on target : %s", targetDir)
200+
}
201+
202+
matched := true
203+
204+
// log.Debug().Msgf("RGS: %v", baseArgs)
205+
206+
cmdJSON := exec.Command(rgPath, baseArgs...)
207+
cmdJSON.Stdout = writer
208+
cmdJSON.Stderr = os.Stderr
209+
210+
log.Debug().Msgf("Creating JSON output for policy %s... ", policy.ID)
211+
err := cmdJSON.Run()
212+
if err != nil {
213+
if exitError, ok := err.(*exec.ExitError); ok {
214+
215+
if exitError.ExitCode() == 1 {
216+
matched = false
217+
err = nil // Reset error as this is the expected outcome for assure
218+
}
219+
220+
if exitError.ExitCode() == 2 {
221+
log.Warn().Msgf("RG exited with code 2")
222+
log.Debug().Msgf("RG Error Args: %v", baseArgs)
223+
if len(exitError.Stderr) > 0 {
224+
log.Debug().Msgf("RG exited with code 2 stderr: %s", string(exitError.Stderr))
225+
}
226+
}
227+
if exitError.ExitCode() != 1 {
228+
log.Error().Err(err).Msg("error executing ripgrep for JSON output")
229+
return matched, fmt.Errorf("error executing ripgrep for JSON output: %w", err)
230+
}
231+
232+
} else {
233+
log.Error().Err(err).Msg("error executing ripgrep for JSON output")
234+
return matched, fmt.Errorf("error executing ripgrep for JSON output: %w", err)
235+
}
236+
}
237+
238+
return matched, nil
239+
}

‎cmd/audit.go

+5-3
Original file line numberDiff line numberDiff line change
@@ -209,9 +209,11 @@ func processPolicy(policy Policy, allFileInfos []FileInfo, rgPath string) {
209209

210210
if policy.Type == "json" || policy.Type == "yaml" || policy.Type == "ini" || policy.Type == "scan" || policy.Type == "assure" {
211211

212-
log.Debug().Msgf(" Processing files for policy %s: ", policy.ID)
213-
for _, file := range filesToProcess {
214-
log.Debug().Msgf(" %s: %s ", file.Path, file.Hash)
212+
log.Debug().Str("policy", policy.ID).Msgf(" Processing files for policy %s ", policy.ID)
213+
if len(filesToProcess) < 15 {
214+
for _, file := range filesToProcess {
215+
log.Debug().Str("policy", policy.ID).Msgf(" %s: %s ", file.Path, file.Hash)
216+
}
215217
}
216218

217219
normalizedID := NormalizeFilename(policy.ID)

‎cmd/aux.go

+30-18
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cmd
22

33
import (
4+
"bufio"
45
"crypto/sha256"
56
"encoding/hex"
67
"encoding/json"
@@ -22,29 +23,40 @@ import (
2223
"gopkg.in/yaml.v3"
2324
)
2425

25-
// patchJSONOutputFile reads the ripgrep JSON output, patches it to create valid JSON, and writes it back to the file
2626
func patchJSONOutputFile(filePath string) error {
27-
// Read the contents of the file
28-
content, err := os.ReadFile(filePath)
27+
// Open the file
28+
file, err := os.Open(filePath)
2929
if err != nil {
30-
return fmt.Errorf("error reading JSON file: %v", err)
30+
return fmt.Errorf("error opening JSON file: %v", err)
3131
}
32+
defer file.Close()
3233

33-
// Split the input into separate JSON objects
34-
objects := strings.Split(string(content), "}\n{")
34+
var validObjects []map[string]interface{}
3535

36-
// Add the missing brackets to create a JSON array
37-
jsonArray := "[" + strings.Join(objects, "},\n{") + "]"
36+
// Read the file line by line
37+
scanner := bufio.NewScanner(file)
38+
for scanner.Scan() {
39+
line := strings.TrimSpace(scanner.Text())
40+
if line == "" {
41+
continue
42+
}
3843

39-
// Parse the JSON array to validate it
40-
var parsed []interface{}
41-
err = json.Unmarshal([]byte(jsonArray), &parsed)
42-
if err != nil {
43-
return fmt.Errorf("error parsing JSON: %v", err)
44+
// Try to parse each line as a separate JSON object
45+
var obj map[string]interface{}
46+
if err := json.Unmarshal([]byte(line), &obj); err == nil {
47+
validObjects = append(validObjects, obj)
48+
} else {
49+
// Log the error and skip this line
50+
fmt.Printf("Skipping invalid JSON line: %s\n", line)
51+
}
52+
}
53+
54+
if err := scanner.Err(); err != nil {
55+
return fmt.Errorf("error reading file: %v", err)
4456
}
4557

46-
// Re-marshal the parsed data to get a properly formatted JSON string
47-
validJSON, err := json.MarshalIndent(parsed, "", " ")
58+
// Marshal the array of objects into a properly formatted JSON string
59+
validJSON, err := json.MarshalIndent(validObjects, "", " ")
4860
if err != nil {
4961
return fmt.Errorf("error marshaling JSON: %v", err)
5062
}
@@ -334,7 +346,7 @@ func createOutputDirectories(isObserve bool) error {
334346
for _, dir := range dirs {
335347
if outputDir != "" {
336348
dir = filepath.Join(outputDir, dir)
337-
log.Debug().Msgf("Creating directory: %s", dir)
349+
// log.Debug().Msgf("Creating directory: %s", dir)
338350

339351
}
340352
if err := os.MkdirAll(dir, 0755); err != nil {
@@ -350,7 +362,7 @@ func cleanupOutputDirectories() error {
350362
if outputDir != "" {
351363
for i, dir := range dirsToClean {
352364
dirsToClean[i] = filepath.Join(outputDir, dir)
353-
log.Debug().Msgf("Cleaning up directories: %v", dirsToClean)
365+
// log.Debug().Msgf("Cleaning up directories: %v", dirsToClean)
354366
}
355367
}
356368
var wg sync.WaitGroup
@@ -367,7 +379,7 @@ func cleanupOutputDirectories() error {
367379
return
368380
}
369381

370-
log.Debug().Msgf("Cleaned up directory: %s ", d)
382+
// log.Debug().Msgf("Cleaned up directory: %s ", d)
371383
}(dir)
372384
}
373385

0 commit comments

Comments
 (0)