Skip to content

Commit

Permalink
Merge pull request #606 from Security-Onion-Solutions/cogburn/ai-desc…
Browse files Browse the repository at this point in the history
…riptions

Cogburn/ai descriptions
  • Loading branch information
coreyogburn authored Aug 8, 2024
2 parents e25cd14 + 342c71c commit 0be0347
Show file tree
Hide file tree
Showing 25 changed files with 1,048 additions and 151 deletions.
7 changes: 4 additions & 3 deletions config/clientparameters.go
Original file line number Diff line number Diff line change
Expand Up @@ -196,9 +196,10 @@ type DetectionsParameters struct {
}

type DetectionParameters struct {
Presets map[string]PresetParameters `json:"presets"`
SeverityTranslations map[string]string `json:"severityTranslations"`
TemplateDetections map[string]string `json:"templateDetections"`
Presets map[string]PresetParameters `json:"presets"`
SeverityTranslations map[string]string `json:"severityTranslations"`
TemplateDetections map[string]string `json:"templateDetections"`
ShowUnreviewedAiSummaries bool `json:"showUnreviewedAiSummaries"`
}

func (params *DetectionsParameters) Verify() error {
Expand Down
6 changes: 5 additions & 1 deletion html/css/app.css
Original file line number Diff line number Diff line change
Expand Up @@ -684,4 +684,8 @@ tbody tr:hover {

.theme--light tbody tr:hover:nth-of-type(even) {
background-color: rgba(230, 230, 230, 0.25) !important;
}
}

.unset-vertical-align {
vertical-align: unset !important;
}
11 changes: 9 additions & 2 deletions html/index.html
Original file line number Diff line number Diff line change
Expand Up @@ -1144,10 +1144,17 @@ <h2 id="detection-title">{{ detect.title }}</h2>
<v-tab-item value="summary" data-aid="detection_summary">
<v-row class="contrast-bg rounded rounded-md pa-4">
<v-col>
<div class="header">{{i18n.summary}}</div>
<div class="summary" data-aid="detection_summary_display">
<div class="header">{{i18n.summary}} <v-icon v-if="showAiSummary()" :title="i18n.aiGenSummary" class="unset-vertical-align" data-aid="detection_summary_ai_indicator">fa-flask</v-icon></div>
<div class="summary" data-aid="detection_ai_summary_display" v-if="showAiSummary()">
{{ detect.aiSummary }}
<div class="font-weight-thin mt-2" v-if="detect.isAiSummaryStale">
({{ i18n.aiSummaryStale }})
</div>
</div>
<div class="summary" v-else data-aid="detection_summary_display">
{{ extractedSummary }}
</div>
</template>
<div class="header">{{i18n.references}}</div>
<div class="reference-links" data-aid="detection_references_display">
<div v-for="r in extractedReferences">
Expand Down
2 changes: 2 additions & 0 deletions html/js/i18n.js
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,8 @@ const i18n = {
address: 'Address',
advanced: 'Temporarily enable advanced interface features',
ago: 'ago',
aiGenSummary: 'AI-Generated Summary',
aiSummaryStale: 'This AI summary was generated for a previous version of the detection and may not be accurate.',
alertAcknowledge: 'Acknowledge',
alertEscalated: 'This alert has already been escalated',
alertUndoAcknowledge: 'Undo Acknowledge',
Expand Down
5 changes: 5 additions & 0 deletions html/js/routes/detection.js
Original file line number Diff line number Diff line change
Expand Up @@ -186,6 +186,7 @@ routes.push({ path: '/detection/:id', name: 'detection', component: {
{ pattern: /condition:/m, message: this.$root.i18n.invalidDetectionStrelkaMissingCondition, match: false },
],
},
showUnreviewedAiSummaries: false,
}},
created() {
this.$root.initializeEditor();
Expand All @@ -211,6 +212,7 @@ routes.push({ path: '/detection/:id', name: 'detection', component: {
this.renderAbbreviatedCount = params["renderAbbreviatedCount"];
this.severityTranslations = params['severityTranslations'];
this.ruleTemplates = params['templateDetections'];
this.showUnreviewedAiSummaries = params['showUnreviewedAiSummaries'];

if (this.$route.params.id === 'create') {
this.detect = this.newDetection();
Expand Down Expand Up @@ -1469,6 +1471,9 @@ routes.push({ path: '/detection/:id', name: 'detection', component: {
},
checkOverrideChangedKey(id, index, key) {
return this.changedOverrideKeys?.[id]?.[index]?.includes(key);
},
showAiSummary() {
return !!(this?.detect?.aiSummary && (this.detect.aiSummaryReviewed || this.showUnreviewedAiSummaries));
}
}
}});
34 changes: 34 additions & 0 deletions html/js/routes/detection.test.js
Original file line number Diff line number Diff line change
Expand Up @@ -1179,4 +1179,38 @@ test('validateSuricata', () => {
comp.detect.publicId = '999999';
msg = comp.validateSuricata();
expect(msg).toBe(null);
});

test('showAiSummary', () => {
comp.detect = null;
expect(comp.showAiSummary()).toBe(false);

comp.detect = { engine: 'strelka' };
expect(comp.showAiSummary()).toBe(false);

comp.detect.aiSummary = 'aiSummary';
expect(comp.showAiSummary()).toBe(false);

comp.detect.aiSummaryReviewed = true;
expect(comp.showAiSummary()).toBe(true);

comp.detect.aiSummary = '';
expect(comp.showAiSummary()).toBe(false);

comp.showUnreviewedAiSummaries = true;

comp.detect = null;
expect(comp.showAiSummary()).toBe(false);

comp.detect = { engine: 'elastalert' };
expect(comp.showAiSummary()).toBe(false);

comp.detect.aiSummary = 'aiSummary';
expect(comp.showAiSummary()).toBe(true);

comp.detect.aiSummaryReviewed = true;
expect(comp.showAiSummary()).toBe(true);

comp.detect.aiSummary = '';
expect(comp.showAiSummary()).toBe(false);
});
16 changes: 16 additions & 0 deletions model/detection.go
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,15 @@ type Detection struct {
// elastalert - sigma only
Product string `json:"product,omitempty"`
Service string `json:"service,omitempty"`

// AI Description fields
*AiFields `json:",omitempty"`
}

type AiFields struct {
AiSummary string `json:"aiSummary"`
AiSummaryReviewed bool `json:"aiSummaryReviewed"`
IsAiSummaryStale bool `json:"isSummaryStale"`
}

type DetectionComment struct {
Expand Down Expand Up @@ -353,3 +362,10 @@ type AuditInfo struct {
Op string
Detection *Detection
}

type AiSummary struct {
PublicId string
Reviewed bool `yaml:"Reviewed"`
Summary string `yaml:"Summary"`
RuleBodyHash string `yaml:"Rule-Body-Hash"`
}
1 change: 1 addition & 0 deletions model/rulerepo.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import (

type RuleRepo struct {
Repo string
Branch *string
License string
Folder *string
Community bool
Expand Down
1 change: 1 addition & 0 deletions server/detectionengine.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ type DetectionEngine interface {
GetState() *model.EngineState
GenerateUnusedPublicId(ctx context.Context) (string, error)
ApplyFilters(detect *model.Detection) (didFilterAct bool, err error)
MergeAuxiliaryData(detect *model.Detection) error
}

type SyncStatus struct {
Expand Down
20 changes: 18 additions & 2 deletions server/detectionhandler.go
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,22 @@ func (h *DetectionHandler) getDetection(w http.ResponseWriter, r *http.Request)
return
}

eng, ok := h.server.DetectionEngines[detect.Engine]
if !ok {
log.WithFields(log.Fields{
"detectionEngine": detect.Engine,
"detectionPublicId": detectId,
}).Error("retrieved detection with unsupported engine")
} else {
err = eng.MergeAuxiliaryData(detect)
if err != nil {
log.WithError(err).WithFields(log.Fields{
"detectionEngine": detect.Engine,
"detectionPublicId": detectId,
}).Error("unable to merge auxiliary data into detection")
}
}

web.Respond(w, r, http.StatusOK, detect)
}

Expand Down Expand Up @@ -597,8 +613,8 @@ func (h *DetectionHandler) bulkUpdateDetectionAsync(ctx context.Context, body *B
filterApplied, err := engine.ApplyFilters(detect)
if err != nil {
logger.WithError(err).WithFields(log.Fields{
"publicId": detect.PublicID,
"engine": detect.Engine,
"detectionPublicId": detect.PublicID,
"detectionEngine": detect.Engine,
}).Error("unable to apply engine filters to detection")

return
Expand Down
146 changes: 146 additions & 0 deletions server/modules/detections/ai_summary.go
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
}
46 changes: 46 additions & 0 deletions server/modules/detections/ai_summary_test.go
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)
}
Loading

0 comments on commit 0be0347

Please sign in to comment.