mongo_utils.go

package utils

import (
	"context"
	"fmt"
	"github.com/rs/zerolog"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

//go:generate go run github.com/maxbrunsfeld/counterfeiter/v6 . MongoUtilsInterface
type MongoUtilsInterface interface {
	GetMongoClient() (*mongo.Client, context.Context, error)
	DisconnectMongoClient(client *mongo.Client, ctx context.Context) error
	GetMongoCollection(client mongo.Client, database, collection string) *mongo.Collection
	CreateCVE(cve *CVE, ctx context.Context, collection *mongo.Collection) error
	GetUniqueDependencies(ctx context.Context, collection *mongo.Collection) ([]CVE, error)
	FindCVE(product string, ctx context.Context, collection *mongo.Collection) (CVE, error)
	RemoveAllCVEs(product string, ctx context.Context, collection *mongo.Collection) error
	CVEExists(cveId string, ctx context.Context, collection *mongo.Collection) (bool, error)
	GetAllCVEs(ctx context.Context, collection *mongo.Collection) ([]CVE, error)
}

type MongoUtils struct {
	DbUri       string
	CtxDuration time.Duration
	Logger      zerolog.Logger
}

type CVE struct {
	ID          primitive.ObjectID `bson:"_id"`
	CreatedAt   time.Time          `bson:"created_at"`
	UpdatedAt   time.Time          `bson:"updated_at"`
	Vendor      string             `bson:"vendor"`
	Product     string             `bson:"product"`
	Severity    string             `bson:"severity"`
	CVEId       string             `bson:"cve_id"`
	Description string             `bson:"description"`
	Score       float64
}

func NewMongoUtils(dbUri string, logger zerolog.Logger) MongoUtils {
	return MongoUtils{
		DbUri:       dbUri,
		CtxDuration: 20 * time.Minute,
		Logger:      logger,
	}
}

func (mu MongoUtils) GetMongoClient() (*mongo.Client, context.Context, error) {
	ctx, _ := context.WithTimeout(context.Background(), mu.CtxDuration)
	clientOptions := options.Client().SetTimeout(30 * time.Second).ApplyURI(mu.DbUri)
	client, err := mongo.Connect(ctx, clientOptions)
	if err != nil {
		return nil, nil, fmt.Errorf("error connecting to MongoDB\n%w", err)
	}

	mu.Logger.Info().Msg("Connected to MongoDB")

	return client, ctx, nil
}

func (mu MongoUtils) GetMongoCollection(client mongo.Client, database, collection string) *mongo.Collection {
	return client.Database(database).Collection(collection)
}

func (mu MongoUtils) DisconnectMongoClient(client *mongo.Client, ctx context.Context) error {
	err := client.Disconnect(ctx)
	if err != nil {
		return fmt.Errorf("error disconnecting client\n%w", err)
	}

	return nil
}

func (mu MongoUtils) CreateCVE(cve *CVE, ctx context.Context, collection *mongo.Collection) error {
	ctx, cancel := context.WithTimeout(ctx, mu.CtxDuration)
	defer cancel()

	_, err := collection.InsertOne(ctx, cve)
	if err != nil {
		return fmt.Errorf("error creating a CVE\n%w", err)
	}

	return nil
}

func (mu MongoUtils) GetUniqueDependencies(ctx context.Context, collection *mongo.Collection) ([]CVE, error) {
	ctx, cancel := context.WithTimeout(ctx, mu.CtxDuration)
	defer cancel()

	results, err := collection.Distinct(ctx, "product", bson.D{{}})

	if err != nil {
		return nil, fmt.Errorf("error getting unique CVEs query result\n%w", err)
	}

	deps := make([]CVE, 0, len(results))
	for _, result := range results {

		cve, err := mu.FindCVE(result.(string), ctx, collection)
		if err != nil {
			return nil, err
		}

		deps = append(deps, cve)
	}

	return deps, nil
}

func (mu MongoUtils) FindCVE(product string, ctx context.Context, collection *mongo.Collection) (CVE, error) {
	ctx, cancel := context.WithTimeout(ctx, mu.CtxDuration)
	defer cancel()

	cve := CVE{}

	result := collection.FindOne(ctx, bson.D{{"product", product}})

	err := result.Decode(&cve)
	if err != nil {
		return CVE{}, err
	}

	return cve, nil
}

func (mu MongoUtils) RemoveAllCVEs(product string, ctx context.Context, collection *mongo.Collection) error {
	ctx, cancel := context.WithTimeout(ctx, mu.CtxDuration)
	defer cancel()

	_, err := collection.DeleteMany(ctx, bson.D{{"product", product}})
	if err != nil {
		return err
	}

	return nil
}

func (mu MongoUtils) CVEExists(cveId string, ctx context.Context, collection *mongo.Collection) (bool, error) {
	ctx, cancel := context.WithTimeout(ctx, mu.CtxDuration)
	defer cancel()

	result, err := collection.CountDocuments(ctx, bson.D{{"cve_id", cveId}})
	if err != nil {
		return false, fmt.Errorf("error retrieving a CVE\n%w", err)
	}

	if result > 0 {
		return true, nil
	}

	return false, nil
}

func (mu MongoUtils) GetAllCVEs(ctx context.Context, collection *mongo.Collection) ([]CVE, error) {
	ctx, cancel := context.WithTimeout(ctx, mu.CtxDuration)
	defer cancel()

	var cves []CVE

	cursor, err := collection.Find(ctx, bson.D{})
	if err != nil {
		return nil, fmt.Errorf("error retrieving all CVEs\n%w", err)
	}

	for cursor.Next(ctx) {
		var cve CVE

		err = cursor.Decode(&cve)
		if err != nil {
			return nil, fmt.Errorf("error decoding CVE from database result\n%w", err)
		}

		cves = append(cves, cve)
	}

	return cves, nil
}