Various bugfixes in Neo4j repository

.env.local

@@ -2,6 +2,8 @@ POSTGRES_USER=postgres
 POSTGRES_PASSWORD=postgres
 POSTGRES_DB=postgres
 SQLITE3_DB=test.db
+NEO4J_AUTH=neo4j/password
+NEO4J_initial_dbms_default__database=amass
 AMASS_DB=example
 AMASS_USER=example
 AMASS_PASSWORD=example

docker-compose.yml

@@ -9,6 +9,20 @@ services:
     ports:
       - "5432:5432"
 
+  neo4j:
+    container_name: assetdb_neo4j
+    image: neo4j:latest
+    restart: always
+    env_file: .env.local
+    expose: 
+      - "7474"
+    ports:
+      - "7474:7474"
+      - "7687:7687"
+    volumes:
+      - neo4j-db:/data    
+
 volumes:
   postgres-db:
     driver: local
+  neo4j-db:

repository/neo4j/db_test.go

@@ -14,11 +14,11 @@ import (
 	neomigrations "github.com/owasp-amass/asset-db/migrations/neo4j"
 )
 
-var store *neoRepository
+var store *NeoRepository
 
 func TestMain(m *testing.M) {
 	var err error
-	dsn := "bolt://neo4j:hackme4fun@localhost:7687/amass"
+	dsn := "bolt://neo4j:password@localhost:7687/amass"
 
 	store, err = New("neo4j", dsn)
 	if err != nil {
@@ -36,7 +36,7 @@ func TestMain(m *testing.M) {
 }
 
 func TestGetDBType(t *testing.T) {
-	if db := store.GetDBType(); db != Neo4j {
+	if db := store.Type(); db != Neo4j {
 		t.Errorf("Failed to return the correct database type")
 	}
 }

repository/neo4j/edge.go

@@ -33,6 +33,13 @@ func (neo *NeoRepository) CreateEdge(ctx context.Context, edge *types.Edge) (*ty
 	if edge.LastSeen.IsZero() {
 		edge.LastSeen = time.Now()
 	}
+
+	if edge.ID != "" {
+		neo.DeleteEdge(ctx, edge.ID)
+	} else {
+		edge.ID = neo.uniqueEntityID()
+	}
+
 	// ensure that duplicate relationships are not entered into the database
 	if e, found := neo.isDuplicateEdge(edge, edge.LastSeen); found {
 		return e, nil
@@ -80,6 +87,9 @@ func (neo *NeoRepository) CreateEdge(ctx context.Context, edge *types.Edge) (*ty
 
 	r.FromEntity = edge.FromEntity
 	r.ToEntity = edge.ToEntity
+
+
+
 	return r, nil
 }
 
@@ -115,7 +125,7 @@ func (neo *NeoRepository) edgeSeen(rel *types.Edge, updated time.Time) error {
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
-	query := fmt.Sprintf("MATCH ()-[r]->() WHERE elementId(r) = $eid SET r.updated_at = localDateTime('%s')", timeToNeo4jTime(updated))
+	query := fmt.Sprintf("MATCH ()-[r {edge_id: $eid}]->() SET r.updated_at = localDateTime('%s')", timeToNeo4jTime(updated))
 	_, err := neo4jdb.ExecuteQuery(ctx, neo.DB, query,
 		map[string]interface{}{
 			"eid": rel.ID,
@@ -132,7 +142,7 @@ func (neo *NeoRepository) FindEdgeById(ctx context.Context, id string) (*types.E
 	defer cancel()
 
 	result, err := neo4jdb.ExecuteQuery(tctx, neo.DB,
-		"MATCH (from:Entity)-[r]->(to:Entity) WHERE elementId(r) = $eid RETURN r, from.entity_id AS fid, to.entity_id AS tid",
+		"MATCH (from:Entity)-[r {edge_id: $eid}]->(to:Entity) RETURN r, from.entity_id AS fid, to.entity_id AS tid",
 		map[string]interface{}{
 			"eid": id,
 		},
@@ -324,7 +334,7 @@ func (neo *NeoRepository) DeleteEdge(ctx context.Context, id string) error {
 	defer cancel()
 
 	_, err := neo4jdb.ExecuteQuery(tctx, neo.DB,
-		"MATCH ()-[r]->() WHERE elementId(r) = $eid DELETE r",
+		"MATCH ()-[r {edge_id: $eid}]->() DELETE r",
 		map[string]interface{}{
 			"eid": id,
 		},

repository/neo4j/edge_tag.go

@@ -62,7 +62,7 @@ func (neo *NeoRepository) CreateEdgeTag(ctx context.Context, edge *types.Edge, i
 		defer cancel()
 
 		result, err := neo4jdb.ExecuteQuery(tctx, neo.DB,
-			"MATCH (n:EdgeTag {tag_id: $tid}) SET p = $props RETURN p",
+			"MATCH (p:EdgeTag {tag_id: $tid}) SET p = $props RETURN p",
 			map[string]interface{}{"tid": tag.ID, "props": props},
 			neo4jdb.EagerResultTransformer,
 			neo4jdb.ExecuteQueryWithDatabase(neo.dbname),

repository/neo4j/entity.go

@@ -41,6 +41,7 @@ func (neo *NeoRepository) CreateEntity(ctx context.Context, input *types.Entity)
 		// ensure that duplicate entities are not entered into the database
 		entity = entities[0]
 		entity.LastSeen = time.Now()
+		entity.Asset = input.Asset
 	}
 
 	if entity != nil {

repository/neo4j/entity_tag.go

@@ -62,7 +62,7 @@ func (neo *NeoRepository) CreateEntityTag(ctx context.Context, entity *types.Ent
 		defer cancel()
 		// update the existing tag
 		result, err := neo4jdb.ExecuteQuery(tctx, neo.DB,
-			"MATCH (n:EntityTag {tag_id: $tid}) SET p = $props RETURN p",
+			"MATCH (p:EntityTag {tag_id: $tid}) SET p = $props RETURN p",
 			map[string]interface{}{"tid": tag.ID, "props": props},
 			neo4jdb.EagerResultTransformer,
 			neo4jdb.ExecuteQueryWithDatabase(neo.dbname),

repository/neo4j/entity_test.go

@@ -11,6 +11,7 @@ import (
 	"net/netip"
 	"testing"
 	"time"
+	"context"
 
 	"github.com/owasp-amass/asset-db/types"
 	oam "github.com/owasp-amass/open-asset-model"
@@ -21,15 +22,15 @@ import (
 )
 
 func TestCreateEntity(t *testing.T) {
-	entity, err := store.CreateEntity(&types.Entity{
+	entity, err := store.CreateEntity(context.TODO(), &types.Entity{
 		Asset: &dns.FQDN{
 			Name: "create1.entity",
 		},
 	})
 	assert.NoError(t, err)
 
 	time.Sleep(250 * time.Millisecond)
-	newer, err := store.CreateEntity(&types.Entity{
+	newer, err := store.CreateEntity(context.TODO(), &types.Entity{
 		Asset: &dns.FQDN{
 			Name: "create1.entity",
 		},
@@ -53,7 +54,7 @@ func TestCreateEntity(t *testing.T) {
 	}
 
 	time.Sleep(250 * time.Millisecond)
-	second, err := store.CreateEntity(&types.Entity{
+	second, err := store.CreateEntity(context.TODO(), &types.Entity{
 		Asset: &dns.FQDN{
 			Name: "create2.entity",
 		},
@@ -67,14 +68,14 @@ func TestCreateEntity(t *testing.T) {
 }
 
 func TestFindEntityById(t *testing.T) {
-	entity, err := store.CreateEntity(&types.Entity{
+	entity, err := store.CreateEntity(context.TODO(), &types.Entity{
 		Asset: &dns.FQDN{
 			Name: "find1.entity",
 		},
 	})
 	assert.NoError(t, err)
 
-	same, err := store.FindEntityById(entity.ID)
+	same, err := store.FindEntityById(context.TODO(), entity.ID)
 	assert.NoError(t, err)
 	assert.Equal(t, entity.ID, same.ID)
 
@@ -90,13 +91,13 @@ func TestFindEntityById(t *testing.T) {
 func TestFindEntitiesByContent(t *testing.T) {
 	fqdn := &dns.FQDN{Name: "findcontent.entity"}
 
-	_, err := store.FindEntitiesByContent(fqdn, time.Time{}, 0)
+	_, err := store.FindEntitiesByContent(context.TODO(), fqdn.AssetType(), time.Time{}, 0, nil)
 	assert.Error(t, err)
 
-	entity, err := store.CreateAsset(fqdn)
+	entity, err := store.CreateAsset(context.TODO(), fqdn)
 	assert.NoError(t, err)
 
-	e, err := store.FindEntitiesByContent(fqdn, entity.CreatedAt.Add(-1*time.Second), 0)
+	e, err := store.FindEntitiesByContent(context.TODO(), fqdn.AssetType(), entity.CreatedAt.Add(-1*time.Second), 0, nil)
 	assert.NoError(t, err)
 	same := e[0]
 	assert.Equal(t, entity.ID, same.ID)
@@ -109,7 +110,7 @@ func TestFindEntitiesByContent(t *testing.T) {
 		t.Errorf("Failed to return an entity with the correct name")
 	}
 
-	_, err = store.FindEntitiesByContent(fqdn, entity.CreatedAt.Add(250*time.Millisecond), 0)
+	_, err = store.FindEntitiesByContent(context.TODO(), fqdn.AssetType(), entity.CreatedAt.Add(250*time.Millisecond), 0, nil)
 	assert.Error(t, err)
 }
 
@@ -120,26 +121,26 @@ func TestFindEntitiesByType(t *testing.T) {
 		addr, err := netip.ParseAddr(fmt.Sprintf("192.168.1.%d", i))
 		assert.NoError(t, err)
 
-		_, err = store.CreateAsset(&oamnet.IPAddress{
+		_, err = store.CreateAsset(context.TODO(), &oamnet.IPAddress{
 			Address: addr,
 			Type:    "IPv4",
 		})
 		assert.NoError(t, err)
 	}
 
-	entities, err := store.FindEntitiesByType(oam.IPAddress, time.Time{}, 0)
+	entities, err := store.FindEntitiesByType(context.TODO(), oam.IPAddress, time.Time{}, 0)
 	assert.NoError(t, err)
 
 	if len(entities) < 10 {
 		t.Errorf("Failed to return the correct number of entities")
 	}
 
 	for i := 1; i <= 10; i++ {
-		_, err := store.CreateAsset(&org.Organization{Name: fmt.Sprintf("findtype%d.entity", i)})
+		_, err := store.CreateAsset(context.TODO(), &org.Organization{Name: fmt.Sprintf("findtype%d.entity", i)})
 		assert.NoError(t, err)
 	}
 
-	entities, err = store.FindEntitiesByType(oam.Organization, now, 0)
+	entities, err = store.FindEntitiesByType(context.TODO(), oam.Organization, now, 0)
 	assert.NoError(t, err)
 
 	if len(entities) < 10 {
@@ -148,16 +149,16 @@ func TestFindEntitiesByType(t *testing.T) {
 }
 
 func TestDeleteEntity(t *testing.T) {
-	entity, err := store.CreateEntity(&types.Entity{
+	entity, err := store.CreateEntity(context.TODO(), &types.Entity{
 		Asset: &dns.FQDN{
 			Name: "delete.entity",
 		},
 	})
 	assert.NoError(t, err)
 
-	err = store.DeleteEntity(entity.ID)
+	err = store.DeleteEntity(context.TODO(), entity.ID)
 	assert.NoError(t, err)
 
-	_, err = store.FindEntityById(entity.ID)
+	_, err = store.FindEntityById(context.TODO(), entity.ID)
 	assert.Error(t, err)
 }

repository/neo4j/extract_edge.go

@@ -16,6 +16,11 @@ import (
 )
 
 func relationshipToEdge(rel neo4jdb.Relationship) (*types.Edge, error) {
+	id, err := neo4jdb.GetProperty[string](rel, "edge_id")
+	if err != nil {
+		return nil, err
+	}
+
 	t, err := neo4jdb.GetProperty[neo4jdb.LocalDateTime](rel, "created_at")
 	if err != nil {
 		return nil, err
@@ -55,7 +60,7 @@ func relationshipToEdge(rel neo4jdb.Relationship) (*types.Edge, error) {
 	}
 
 	return &types.Edge{
-		ID:        rel.GetElementId(),
+		ID:        id,
 		CreatedAt: created,
 		LastSeen:  updated,
 		Relation:  r,

repository/neo4j/query_relations.go

@@ -22,6 +22,7 @@ func edgePropsMap(edge *types.Edge) (map[string]interface{}, error) {
 
 	m := make(map[string]interface{})
 	// begin populating the map of parameters
+	m["edge_id"] = edge.ID
 	m["etype"] = edge.Relation.RelationType()
 	m["created_at"] = timeToNeo4jTime(edge.CreatedAt)
 	m["updated_at"] = timeToNeo4jTime(edge.LastSeen)

repository/neo4j/tag_test.go

@@ -89,13 +89,16 @@ func TestEntityTag(t *testing.T) {
 }
 
 func TestEdgeTag(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
 	e1, err := store.CreateAsset(ctx, &dns.FQDN{Name: "owasp.org"})
 	assert.NoError(t, err)
 
 	e2, err := store.CreateAsset(ctx, &dns.FQDN{Name: "www.owasp.org"})
 	assert.NoError(t, err)
 
-	edge, err := store.CreateEdge(&types.Edge{
+	edge, err := store.CreateEdge(ctx, &types.Edge{
 		Relation: &dns.BasicDNSRelation{
 			Name:   "dns_record",
 			Header: dns.RRHeader{RRType: 5},