Add global settings (#40)

database/db.go

@@ -25,6 +25,7 @@ const (
 	UserCollectionName       = "users"
 	TicketTypeCollectionName = "tickettypes"
 	JobCollectionName        = "jobs"
+	SettingsCollectionName   = "settings"
 
 	TicketArtifactsGraphName     = "Graph"
 	RelatedTicketsCollectionName = "related"
@@ -44,6 +45,7 @@ type Database struct {
 	userCollection       *busdb.Collection
 	tickettypeCollection *busdb.Collection
 	jobCollection        *busdb.Collection
+	settingsCollection   *busdb.Collection
 
 	relatedCollection *busdb.Collection
 	// containsCollection *busdb.Collection
@@ -122,6 +124,10 @@ func New(ctx context.Context, index *index.Index, bus *bus.Bus, hooks *hooks.Hoo
 	if err != nil {
 		return nil, err
 	}
+	settingsCollection, err := arangoDB.Collection(ctx, SettingsCollectionName)
+	if err != nil {
+		return nil, err
+	}
 
 	hookedDB, err := busdb.NewDatabase(ctx, arangoDB, bus)
 	if err != nil {
@@ -142,6 +148,7 @@ func New(ctx context.Context, index *index.Index, bus *bus.Bus, hooks *hooks.Hoo
 		userCollection:       busdb.NewCollection(userCollection, hookedDB),
 		tickettypeCollection: busdb.NewCollection(tickettypeCollection, hookedDB),
 		jobCollection:        busdb.NewCollection(jobCollection, hookedDB),
+		settingsCollection:   busdb.NewCollection(settingsCollection, hookedDB),
 	}
 
 	return db, nil
@@ -189,5 +196,6 @@ func (db *Database) Truncate(ctx context.Context) {
 	db.tickettypeCollection.Truncate(ctx)
 	db.jobCollection.Truncate(ctx)
 	db.relatedCollection.Truncate(ctx)
+	db.settingsCollection.Truncate(ctx)
 	// db.containsCollection.Truncate(ctx)
 }

database/migrations/migrations.go

@@ -51,6 +51,11 @@ func generateMigrations() ([]Migration, error) {
 		&createCollection{ID: "create-job-collection", Name: "jobs", DataType: "job", Schema: `{"properties":{"automation":{"type":"string"},"log":{"type":"string"},"payload":{},"origin":{"properties":{"artifact_origin":{"properties":{"artifact":{"type":"string"},"ticket_id":{"format":"int64","type":"integer"}},"required":["artifact","ticket_id"],"type":"object"},"task_origin":{"properties":{"playbook_id":{"type":"string"},"task_id":{"type":"string"},"ticket_id":{"format":"int64","type":"integer"}},"required":["playbook_id","task_id","ticket_id"],"type":"object"}},"type":"object"},"output":{"properties":{},"type":"object"},"running":{"type":"boolean"},"status":{"type":"string"}},"required":["automation","running","status"],"type":"object"}`},
 
 		&createDocument{ID: "create-playbook-simple", Collection: "playbooks", Document: &busdb.Keyed{Key: "simple", Doc: model.PlaybookTemplate{Name: "Simple", Yaml: SimplePlaybook}}},
+
+		&createCollection{ID: "create-settings-collection", Name: "settings", DataType: "settings", Schema: `{"type":"object","properties":{"artifactStates":{"title":"Artifact States","items":{"type":"object","properties":{"color":{"title":"Color","type":"string","enum":["error","info","success","warning"]},"icon":{"title":"Icon (https://materialdesignicons.com)","type":"string"},"id":{"title":"ID","type":"string"},"name":{"title":"Name","type":"string"}},"required":["id","name","icon"]},"type":"array"},"artifactKinds":{"title":"Artifact Kinds","items":{"type":"object","properties":{"color":{"title":"Color","type":"string","enum":["error","info","success","warning"]},"icon":{"title":"Icon (https://materialdesignicons.com)","type":"string"},"id":{"title":"ID","type":"string"},"name":{"title":"Name","type":"string"}},"required":["id","name","icon"]},"type":"array"},"timeformat":{"title":"Time Format","type":"string"}},"required":["timeformat","artifactKinds","artifactStates"]}`},
+		&createDocument{ID: "create-settings-global", Collection: "settings", Document: &busdb.Keyed{Key: "global", Doc: model.Settings{ArtifactStates: []*model.Type{{Icon: "mdi-help-circle-outline", ID: "unknown", Name: "Unknown", Color: pointer.String(model.TypeColorInfo)}, {Icon: "mdi-skull", ID: "malicious", Name: "Malicious", Color: pointer.String(model.TypeColorError)}, {Icon: "mdi-check", ID: "clean", Name: "Clean", Color: pointer.String(model.TypeColorSuccess)}}, ArtifactKinds: []*model.Type{{Icon: "mdi-server", ID: "asset", Name: "Asset"}, {Icon: "mdi-bullseye", ID: "ioc", Name: "IOC"}}, Timeformat: "YYYY-MM-DDThh:mm:ss"}}},
+
+		&updateSchema{ID: "update-ticket-collection", Name: "tickets", DataType: "ticket", Schema: `{"properties":{"artifacts":{"items":{"properties":{"enrichments":{"additionalProperties":{"properties":{"created":{"format":"date-time","type":"string"},"data":{"example":{"hash":"b7a067a742c20d07a7456646de89bc2d408a1153"},"properties":{},"type":"object"},"name":{"example":"hash.sha1","type":"string"}},"required":["created","data","name"],"type":"object"},"type":"object"},"name":{"example":"2.2.2.2","type":"string"},"status":{"example":"Unknown","type":"string"},"type":{"type":"string"},"kind":{"type":"string"}},"required":["name"],"type":"object"},"type":"array"},"comments":{"items":{"properties":{"created":{"format":"date-time","type":"string"},"creator":{"type":"string"},"message":{"type":"string"}},"required":["created","creator","message"],"type":"object"},"type":"array"},"created":{"format":"date-time","type":"string"},"details":{"example":{"description":"my little incident"},"properties":{},"type":"object"},"files":{"items":{"properties":{"key":{"example":"myfile","type":"string"},"name":{"example":"notes.docx","type":"string"}},"required":["key","name"],"type":"object"},"type":"array"},"modified":{"format":"date-time","type":"string"},"name":{"example":"WannyCry","type":"string"},"owner":{"example":"bob","type":"string"},"playbooks":{"additionalProperties":{"properties":{"name":{"example":"Phishing","type":"string"},"tasks":{"additionalProperties":{"properties":{"automation":{"type":"string"},"closed":{"format":"date-time","type":"string"},"created":{"format":"date-time","type":"string"},"data":{"properties":{},"type":"object"},"done":{"type":"boolean"},"join":{"example":false,"type":"boolean"},"payload":{"additionalProperties":{"type":"string"},"type":"object"},"name":{"example":"Inform user","type":"string"},"next":{"additionalProperties":{"type":"string"},"type":"object"},"owner":{"type":"string"},"schema":{"properties":{},"type":"object"},"type":{"enum":["task","input","automation"],"example":"task","type":"string"}},"required":["created","done","name","type"],"type":"object"},"type":"object"}},"required":["name","tasks"],"type":"object"},"type":"object"},"read":{"example":["bob"],"items":{"type":"string"},"type":"array"},"references":{"items":{"properties":{"href":{"example":"https://cve.mitre.org/cgi-bin/cvename.cgi?name=cve-2017-0144","type":"string"},"name":{"example":"CVE-2017-0144","type":"string"}},"required":["href","name"],"type":"object"},"type":"array"},"schema":{"example":"{}","type":"string"},"status":{"example":"open","type":"string"},"type":{"example":"incident","type":"string"},"write":{"example":["alice"],"items":{"type":"string"},"type":"array"}},"required":["created","modified","name","schema","status","type"],"type":"object"}`},
 	}, nil
 }
 

database/settings.go

@@ -2,84 +2,34 @@ package database
 
 import (
 	"context"
-	"errors"
 
-	"github.com/arangodb/go-driver"
-
-	"github.com/SecurityBrewery/catalyst/database/busdb"
 	"github.com/SecurityBrewery/catalyst/generated/model"
 )
 
-func toUserDataResponse(key string, doc *model.UserData) *model.UserDataResponse {
-	return &model.UserDataResponse{
-		Email:      doc.Email,
-		ID:         key,
-		Image:      doc.Image,
-		Name:       doc.Name,
-		Timeformat: doc.Timeformat,
+func (db *Database) Settings(ctx context.Context) (*model.Settings, error) {
+	settings := &model.Settings{}
+	if _, err := db.settingsCollection.ReadDocument(ctx, "global", settings); err != nil {
+		return nil, err
 	}
+
+	return settings, nil
 }
 
-func (db *Database) UserDataCreate(ctx context.Context, id string, userdata *model.UserData) error {
-	if userdata == nil {
-		return errors.New("requires setting")
-	}
-	if id == "" {
-		return errors.New("requires username")
-	}
-
-	_, err := db.userdataCollection.CreateDocument(ctx, ctx, id, userdata)
-	return err
-}
-
-func (db *Database) UserDataGetOrCreate(ctx context.Context, id string, newUserData *model.UserData) (*model.UserDataResponse, error) {
-	setting, err := db.UserDataGet(ctx, id)
-	if err != nil {
-		return toUserDataResponse(id, newUserData), db.UserDataCreate(ctx, id, newUserData)
-	}
-	return setting, nil
-}
-
-func (db *Database) UserDataGet(ctx context.Context, id string) (*model.UserDataResponse, error) {
-	var doc model.UserData
-	meta, err := db.userdataCollection.ReadDocument(ctx, id, &doc)
+func (db *Database) SaveSettings(ctx context.Context, settings *model.Settings) (*model.Settings, error) {
+	exists, err := db.settingsCollection.DocumentExists(ctx, "global")
 	if err != nil {
 		return nil, err
 	}
 
-	return toUserDataResponse(meta.Key, &doc), err
-}
-
-func (db *Database) UserDataList(ctx context.Context) ([]*model.UserDataResponse, error) {
-	query := "FOR d IN @@collection SORT d.username ASC RETURN d"
-	cursor, _, err := db.Query(ctx, query, map[string]interface{}{"@collection": UserDataCollectionName}, busdb.ReadOperation)
-	if err != nil {
-		return nil, err
-	}
-	defer cursor.Close()
-	var docs []*model.UserDataResponse
-	for {
-		var doc model.UserData
-		meta, err := cursor.ReadDocument(ctx, &doc)
-		if driver.IsNoMoreDocuments(err) {
-			break
-		} else if err != nil {
+	if exists {
+		if _, err := db.settingsCollection.ReplaceDocument(ctx, "global", settings); err != nil {
+			return nil, err
+		}
+	} else {
+		if _, err := db.settingsCollection.CreateDocument(ctx, ctx, "global", settings); err != nil {
 			return nil, err
 		}
-		docs = append(docs, toUserDataResponse(meta.Key, &doc))
 	}
 
-	return docs, err
-}
-
-func (db *Database) UserDataUpdate(ctx context.Context, id string, userdata *model.UserData) (*model.UserDataResponse, error) {
-	var doc model.UserData
-	ctx = driver.WithReturnNew(ctx, &doc)
-
-	meta, err := db.userdataCollection.ReplaceDocument(ctx, id, userdata)
-	if err != nil {
-		return nil, err
-	}
-
-	return toUserDataResponse(meta.Key, &doc), nil
+	return settings, nil
 }

database/userdata.go

@@ -0,0 +1,85 @@
+package database
+
+import (
+	"context"
+	"errors"
+
+	"github.com/arangodb/go-driver"
+
+	"github.com/SecurityBrewery/catalyst/database/busdb"
+	"github.com/SecurityBrewery/catalyst/generated/model"
+)
+
+func toUserDataResponse(key string, doc *model.UserData) *model.UserDataResponse {
+	return &model.UserDataResponse{
+		Email:      doc.Email,
+		ID:         key,
+		Image:      doc.Image,
+		Name:       doc.Name,
+		Timeformat: doc.Timeformat,
+	}
+}
+
+func (db *Database) UserDataCreate(ctx context.Context, id string, userdata *model.UserData) error {
+	if userdata == nil {
+		return errors.New("requires setting")
+	}
+	if id == "" {
+		return errors.New("requires username")
+	}
+
+	_, err := db.userdataCollection.CreateDocument(ctx, ctx, id, userdata)
+	return err
+}
+
+func (db *Database) UserDataGetOrCreate(ctx context.Context, id string, newUserData *model.UserData) (*model.UserDataResponse, error) {
+	setting, err := db.UserDataGet(ctx, id)
+	if err != nil {
+		return toUserDataResponse(id, newUserData), db.UserDataCreate(ctx, id, newUserData)
+	}
+	return setting, nil
+}
+
+func (db *Database) UserDataGet(ctx context.Context, id string) (*model.UserDataResponse, error) {
+	var doc model.UserData
+	meta, err := db.userdataCollection.ReadDocument(ctx, id, &doc)
+	if err != nil {
+		return nil, err
+	}
+
+	return toUserDataResponse(meta.Key, &doc), err
+}
+
+func (db *Database) UserDataList(ctx context.Context) ([]*model.UserDataResponse, error) {
+	query := "FOR d IN @@collection SORT d.username ASC RETURN d"
+	cursor, _, err := db.Query(ctx, query, map[string]interface{}{"@collection": UserDataCollectionName}, busdb.ReadOperation)
+	if err != nil {
+		return nil, err
+	}
+	defer cursor.Close()
+	var docs []*model.UserDataResponse
+	for {
+		var doc model.UserData
+		meta, err := cursor.ReadDocument(ctx, &doc)
+		if driver.IsNoMoreDocuments(err) {
+			break
+		} else if err != nil {
+			return nil, err
+		}
+		docs = append(docs, toUserDataResponse(meta.Key, &doc))
+	}
+
+	return docs, err
+}
+
+func (db *Database) UserDataUpdate(ctx context.Context, id string, userdata *model.UserData) (*model.UserDataResponse, error) {
+	var doc model.UserData
+	ctx = driver.WithReturnNew(ctx, &doc)
+
+	meta, err := db.userdataCollection.ReplaceDocument(ctx, id, userdata)
+	if err != nil {
+		return nil, err
+	}
+
+	return toUserDataResponse(meta.Key, &doc), nil
+}