-
Notifications
You must be signed in to change notification settings - Fork 20
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #606 from Security-Onion-Solutions/cogburn/ai-desc…
…riptions Cogburn/ai descriptions
- Loading branch information
Showing
25 changed files
with
1,048 additions
and
151 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,146 @@ | ||
package detections | ||
|
||
import ( | ||
"errors" | ||
"fmt" | ||
"net/url" | ||
"path" | ||
"path/filepath" | ||
"sync" | ||
"time" | ||
|
||
"github.com/security-onion-solutions/securityonion-soc/model" | ||
|
||
"github.com/apex/log" | ||
"gopkg.in/yaml.v3" | ||
) | ||
|
||
var aiRepoMutex = sync.RWMutex{} | ||
var lastSuccessfulAiUpdate time.Time | ||
|
||
type AiLoader interface { | ||
LoadAuxiliaryData(summaries []*model.AiSummary) error | ||
} | ||
|
||
//go:generate mockgen -destination mock/mock_ailoader.go -package mock . AiLoader | ||
|
||
func RefreshAiSummaries(eng AiLoader, lang model.SigLanguage, isRunning *bool, aiRepoPath string, aiRepoUrl string, aiRepoBranch string, logger *log.Entry, iom IOManager) error { | ||
err := updateAiRepo(isRunning, aiRepoPath, aiRepoUrl, aiRepoBranch, iom) | ||
if err != nil { | ||
if errors.Is(err, ErrModuleStopped) { | ||
return err | ||
} | ||
|
||
logger.WithError(err).WithFields(log.Fields{ | ||
"aiRepoUrl": aiRepoUrl, | ||
"aiRepoPath": aiRepoPath, | ||
}).Error("unable to update AI repo") | ||
} | ||
|
||
parser, err := url.Parse(aiRepoUrl) | ||
if err != nil { | ||
log.WithError(err).WithField("aiRepoUrl", aiRepoUrl).Error("failed to parse repo URL, doing nothing with it") | ||
} else { | ||
_, lastFolder := path.Split(parser.Path) | ||
repoPath := filepath.Join(aiRepoPath, lastFolder) | ||
|
||
sums, err := readAiSummary(isRunning, repoPath, lang, logger, iom) | ||
if err != nil { | ||
logger.WithError(err).WithField("repoPath", repoPath).Error("unable to read AI summaries") | ||
} else { | ||
err = eng.LoadAuxiliaryData(sums) | ||
if err != nil { | ||
logger.WithError(err).Error("unable to load AI summaries") | ||
} else { | ||
logger.Info("successfully loaded AI summaries") | ||
} | ||
} | ||
} | ||
|
||
return nil | ||
} | ||
|
||
func updateAiRepo(isRunning *bool, baseRepoFolder string, repoUrl string, branch string, iom IOManager) error { | ||
if time.Since(lastSuccessfulAiUpdate) < time.Second*5 { | ||
log.Info("AI summary repo was updated recently, skipping update") | ||
return nil | ||
} | ||
|
||
aiRepoMutex.Lock() | ||
defer aiRepoMutex.Unlock() | ||
|
||
if time.Since(lastSuccessfulAiUpdate) < time.Second*5 { | ||
log.Info("AI summary repo was updated recently, skipping update") | ||
return nil | ||
} | ||
|
||
var branchPtr *string | ||
if branch != "" { | ||
branchPtr = &branch | ||
} | ||
|
||
_, _, err := UpdateRepos(isRunning, baseRepoFolder, []*model.RuleRepo{ | ||
{ | ||
Repo: repoUrl, | ||
Branch: branchPtr, | ||
}, | ||
}, iom) | ||
|
||
if err == nil { | ||
lastSuccessfulAiUpdate = time.Now() | ||
} | ||
|
||
return err | ||
} | ||
|
||
func readAiSummary(isRunning *bool, repoRoot string, lang model.SigLanguage, logger *log.Entry, iom IOManager) (sums []*model.AiSummary, err error) { | ||
aiRepoMutex.RLock() | ||
defer aiRepoMutex.RUnlock() | ||
|
||
filename := fmt.Sprintf("%s_summaries.yaml", lang) | ||
targetFile := filepath.Join(repoRoot, "detections-ai/", filename) | ||
|
||
logger.WithField("targetFile", targetFile).Info("reading AI summaries") | ||
|
||
raw, err := iom.ReadFile(targetFile) | ||
if err != nil { | ||
return nil, err | ||
} | ||
|
||
// large yaml files take 30+ seconds to unmarshal, so we need to check if the | ||
// module has stopped or risk becoming unresponsive when sent a signal to stop | ||
done := false | ||
data := map[string]*model.AiSummary{} | ||
|
||
go func() { | ||
err = yaml.Unmarshal(raw, data) | ||
done = true | ||
}() | ||
|
||
for !done { | ||
if !*isRunning { | ||
return nil, ErrModuleStopped | ||
} | ||
|
||
time.Sleep(time.Millisecond * 200) | ||
} | ||
|
||
if err != nil { | ||
return nil, err | ||
} | ||
|
||
logger.Info("successfully unmarshalled AI summaries, parsing...") | ||
|
||
for pid, sum := range data { | ||
if !*isRunning { | ||
return nil, ErrModuleStopped | ||
} | ||
|
||
sum.PublicId = pid | ||
sums = append(sums, sum) | ||
} | ||
|
||
logger.WithField("aiSummaryCount", len(sums)).Info("successfully parsed AI summaries") | ||
|
||
return sums, nil | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,46 @@ | ||
package detections | ||
|
||
import ( | ||
"io/fs" | ||
"testing" | ||
|
||
"github.com/apex/log" | ||
"github.com/security-onion-solutions/securityonion-soc/model" | ||
"github.com/security-onion-solutions/securityonion-soc/server/modules/detections/mock" | ||
|
||
"github.com/tj/assert" | ||
"go.uber.org/mock/gomock" | ||
) | ||
|
||
func TestRefreshAiSummaries(t *testing.T) { | ||
ctrl := gomock.NewController(t) | ||
defer ctrl.Finish() | ||
|
||
isRunning := true | ||
repo := "http://github.com/user/repo1" | ||
branch := "generated-summaries-stable" | ||
summaries := `{"87e55c67-46f0-4a7b-a3c6-d473ab7e8392": { "Reviewed": false, "Summary": "ai text goes here"}, "a23077fc-a5ef-427f-92ab-d3de7f56834d": { "Reviewed": true, "Summary": "ai text goes here" } }` | ||
|
||
iom := mock.NewMockIOManager(ctrl) | ||
loader := mock.NewMockAiLoader(ctrl) | ||
|
||
iom.EXPECT().ReadDir("baseRepoFolder").Return([]fs.DirEntry{}, nil) | ||
iom.EXPECT().CloneRepo(gomock.Any(), "baseRepoFolder/repo1", repo, &branch).Return(nil) | ||
iom.EXPECT().ReadFile("baseRepoFolder/repo1/detections-ai/sigma_summaries.yaml").Return([]byte(summaries), nil) | ||
loader.EXPECT().LoadAuxiliaryData([]*model.AiSummary{ | ||
{ | ||
PublicId: "87e55c67-46f0-4a7b-a3c6-d473ab7e8392", | ||
Summary: "ai text goes here", | ||
}, | ||
{ | ||
PublicId: "a23077fc-a5ef-427f-92ab-d3de7f56834d", | ||
Reviewed: true, | ||
Summary: "ai text goes here", | ||
}, | ||
}).Return(nil) | ||
|
||
logger := log.WithField("test", true) | ||
|
||
err := RefreshAiSummaries(loader, model.SigLangSigma, &isRunning, "baseRepoFolder", repo, branch, logger, iom) | ||
assert.NoError(t, err) | ||
} |
Oops, something went wrong.