From 3e0777b89811f60eed3f02052c22bb5a152c62e7 Mon Sep 17 00:00:00 2001 From: JOppliger Date: Tue, 7 Apr 2026 12:04:52 +0200 Subject: [PATCH] Various bugfixes in Neo4j repository --- .env.local | 2 ++ docker-compose.yml | 14 ++++++++++++ repository/neo4j/db_test.go | 6 +++--- repository/neo4j/edge.go | 16 +++++++++++--- repository/neo4j/edge_tag.go | 2 +- repository/neo4j/entity.go | 1 + repository/neo4j/entity_tag.go | 2 +- repository/neo4j/entity_test.go | 33 +++++++++++++++-------------- repository/neo4j/extract_edge.go | 7 +++++- repository/neo4j/query_relations.go | 1 + repository/neo4j/tag_test.go | 5 ++++- 11 files changed, 63 insertions(+), 26 deletions(-) diff --git a/.env.local b/.env.local index 03a3282..097deb4 100644 --- a/.env.local +++ b/.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 \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml index 51f87ff..82bf570 100644 --- a/docker-compose.yml +++ b/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: diff --git a/repository/neo4j/db_test.go b/repository/neo4j/db_test.go index ecc1ddb..4027298 100644 --- a/repository/neo4j/db_test.go +++ b/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") } } diff --git a/repository/neo4j/edge.go b/repository/neo4j/edge.go index 387aa3e..2ecfa09 100644 --- a/repository/neo4j/edge.go +++ b/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, }, diff --git a/repository/neo4j/edge_tag.go b/repository/neo4j/edge_tag.go index 73cd437..bd4e010 100644 --- a/repository/neo4j/edge_tag.go +++ b/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), diff --git a/repository/neo4j/entity.go b/repository/neo4j/entity.go index 34a22e0..c01989f 100644 --- a/repository/neo4j/entity.go +++ b/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 { diff --git a/repository/neo4j/entity_tag.go b/repository/neo4j/entity_tag.go index 19c8fc9..cfc66d4 100644 --- a/repository/neo4j/entity_tag.go +++ b/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), diff --git a/repository/neo4j/entity_test.go b/repository/neo4j/entity_test.go index b508202..b56fa83 100644 --- a/repository/neo4j/entity_test.go +++ b/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,7 +22,7 @@ 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", }, @@ -29,7 +30,7 @@ func TestCreateEntity(t *testing.T) { 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,14 +121,14 @@ 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 { @@ -135,11 +136,11 @@ func TestFindEntitiesByType(t *testing.T) { } 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) } diff --git a/repository/neo4j/extract_edge.go b/repository/neo4j/extract_edge.go index 62e3988..ef704fa 100644 --- a/repository/neo4j/extract_edge.go +++ b/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, diff --git a/repository/neo4j/query_relations.go b/repository/neo4j/query_relations.go index 3ed8e3a..da3de4f 100644 --- a/repository/neo4j/query_relations.go +++ b/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) diff --git a/repository/neo4j/tag_test.go b/repository/neo4j/tag_test.go index a1ea529..e791f87 100644 --- a/repository/neo4j/tag_test.go +++ b/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},