Initial commit

1 files changed, 137 insertions, 0 deletions

diff --git a/storage.go b/storage.go
new file mode 100644
index 0000000..e8f6645
--- /dev/null
+++ b/storage.go
@@ -0,0 +1,137 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+)
+
+const defaultConnectionTimeout = 10*time.Second
+const defaultOperationTimeout = 3*time.Second
+
+type Storage struct {
+	client *mongo.Client
+	collections map[string]*mongo.Collection
+}
+
+type OrderedDocument = bson.D
+type UnorderedDocument = bson.M
+
+func NewStorage(uri string, port int, database string) Storage {
+	opt := options.Client()
+	opt.ApplyURI(fmt.Sprintf("mongodb://%s:%v", uri, port))
+	client, err := mongo.NewClient(opt)
+	if err != nil {
+		panic("Failed to create mongo client")
+	}
+
+	db := client.Database(database)
+	colls := map[string]*mongo.Collection{
+		"imported_pcaps": db.Collection("imported_pcaps"),
+		"connections": db.Collection("connections"),
+	}
+
+	return Storage{
+		client:      client,
+		collections: colls,
+	}
+}
+
+func (storage *Storage) Connect(ctx context.Context) error {
+	if ctx == nil {
+		ctx, _ = context.WithTimeout(context.Background(), defaultConnectionTimeout)
+	}
+
+	return storage.client.Connect(ctx)
+}
+
+func (storage *Storage) InsertOne(ctx context.Context, collectionName string,
+	document interface{}) (interface{}, error) {
+
+	collection, ok := storage.collections[collectionName]
+	if !ok {
+		return nil, errors.New("invalid collection: " + collectionName)
+	}
+
+	if ctx == nil {
+		ctx, _ = context.WithTimeout(context.Background(), defaultOperationTimeout)
+	}
+
+	result, err := collection.InsertOne(ctx, document)
+	if err != nil {
+		return nil, err
+	}
+
+	return result.InsertedID, nil
+}
+
+func (storage *Storage) UpdateOne(ctx context.Context, collectionName string,
+	filter interface{}, update interface {}, upsert bool) (interface{}, error) {
+
+	collection, ok := storage.collections[collectionName]
+	if !ok {
+		return nil, errors.New("invalid collection: " + collectionName)
+	}
+
+	if ctx == nil {
+		ctx, _ = context.WithTimeout(context.Background(), defaultOperationTimeout)
+	}
+
+	opts := options.Update().SetUpsert(upsert)
+	update = bson.D{{"$set", update}}
+
+	result, err := collection.UpdateOne(ctx, filter, update, opts)
+	if err != nil {
+		return nil, err
+	}
+
+	if upsert {
+		return result.UpsertedID, nil
+	}
+
+	return result.ModifiedCount == 1, nil
+}
+
+func (storage *Storage) FindOne(ctx context.Context, collectionName string,
+	filter interface{}) (UnorderedDocument, error) {
+
+	collection, ok := storage.collections[collectionName]
+	if !ok {
+		return nil, errors.New("invalid collection: " + collectionName)
+	}
+
+	if ctx == nil {
+		ctx, _ = context.WithTimeout(context.Background(), defaultOperationTimeout)
+	}
+
+	var result bson.M
+	err := collection.FindOne(ctx, filter).Decode(&result)
+	if err != nil {
+		if err == mongo.ErrNoDocuments {
+			return nil, nil
+		}
+
+		return nil, err
+	}
+
+	return result, nil
+}
+
+
+func testStorage() {
+	storage := NewStorage("localhost", 27017, "testing")
+	_ = storage.Connect(nil)
+
+	id, err := storage.InsertOne(nil, "connections", bson.M{"_id": "provaaa"})
+	if err != nil {
+		panic(err)
+	} else {
+		fmt.Println(id)
+	}
+
+}