From 5ade17b9c0cb1e2a3d46a6d68b000adc009ff71e Mon Sep 17 00:00:00 2001
From: Logan Lembke <logo2462@gmail.com>
Date: Tue, 13 Apr 2021 10:12:32 -0600
Subject: [PATCH] Add Indices to Quickly Search for Hosts which Contacted BL
 Hosts (#627)

* Add blacklisted host indices

* Fix method missing from interface

Co-authored-by: Logan L <logan@bhis.co>
---
 parser/fsimporter.go        |  5 +++++
 pkg/blacklist/mongodb.go    | 23 +++++++++++++++++++++++
 pkg/blacklist/repository.go |  1 +
 3 files changed, 29 insertions(+)

diff --git a/parser/fsimporter.go b/parser/fsimporter.go
index 9929d731..4b282174 100644
--- a/parser/fsimporter.go
+++ b/parser/fsimporter.go
@@ -892,6 +892,11 @@ func (fs *FSImporter) markBlacklistedPeers(hostMap map[string]*host.Input) {
 	if len(hostMap) > 0 {
 		blacklistRepo := blacklist.NewMongoRepository(fs.res)
 
+		err := blacklistRepo.CreateIndexes()
+		if err != nil {
+			fs.res.Log.Error(err)
+		}
+
 		// send uconns to host analysis
 		blacklistRepo.Upsert()
 	}
diff --git a/pkg/blacklist/mongodb.go b/pkg/blacklist/mongodb.go
index 0e1232d8..553653c7 100644
--- a/pkg/blacklist/mongodb.go
+++ b/pkg/blacklist/mongodb.go
@@ -7,6 +7,7 @@ import (
 	"github.com/activecm/rita/pkg/data"
 	"github.com/activecm/rita/resources"
 	"github.com/activecm/rita/util"
+	"github.com/globalsign/mgo"
 	"github.com/globalsign/mgo/bson"
 )
 
@@ -21,6 +22,28 @@ func NewMongoRepository(res *resources.Resources) Repository {
 	}
 }
 
+//CreateIndexes sets up the indices needed to find hosts which contacted blacklisted hosts
+func (r *repo) CreateIndexes() error {
+	session := r.res.DB.Session.Copy()
+	defer session.Close()
+
+	coll := session.DB(r.res.DB.GetSelectedDB()).C(r.res.Config.T.Structure.HostTable)
+
+	// create hosts collection
+	// Desired indexes
+	indexes := []mgo.Index{
+		{Key: []string{"dat.bl.ip", "dat.bl.network_uuid"}},
+	}
+
+	for _, index := range indexes {
+		err := coll.EnsureIndex(index)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 //Upsert loops through every domain ....
 func (r *repo) Upsert() {
 
diff --git a/pkg/blacklist/repository.go b/pkg/blacklist/repository.go
index 95939d4d..522de39f 100644
--- a/pkg/blacklist/repository.go
+++ b/pkg/blacklist/repository.go
@@ -7,6 +7,7 @@ import (
 
 // Repository for blacklist results in host collection
 type Repository interface {
+	CreateIndexes() error
 	Upsert()
 }