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 <[email protected]>
activecm · Apr 13, 2021 · 5ade17b · 5ade17b
1 parent ce2efdc
commit 5ade17b
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 0 deletions.
diff --git 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
@@ -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
@@ -7,6 +7,7 @@ import (
 
 // Repository for blacklist results in host collection
 type Repository interface {
+	CreateIndexes() error
 	Upsert()
 }