Cogburn/ai descriptions

config/clientparameters.go

@@ -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 {

html/css/app.css

@@ -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;
+}

html/index.html

@@ -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">

html/js/i18n.js

@@ -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',

html/js/routes/detection.js

@@ -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();
@@ -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();
@@ -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));
 		}
 	}
 }});

html/js/routes/detection.test.js

@@ -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);
 });

model/detection.go

@@ -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 {
@@ -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"`
+}

model/rulerepo.go

@@ -12,6 +12,7 @@ import (
 
 type RuleRepo struct {
 	Repo      string
+	Branch    *string
 	License   string
 	Folder    *string
 	Community bool

server/detectionengine.go

@@ -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 {

server/detectionhandler.go

@@ -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)
 }
 
@@ -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

server/modules/detections/ai_summary.go

@@ -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
+}

server/modules/detections/ai_summary_test.go

@@ -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)
+}