From a88762219f52d16b4edd9d4156e7cca3cadd6049 Mon Sep 17 00:00:00 2001 From: caffix Date: Sat, 21 Dec 2024 12:38:12 -0500 Subject: [PATCH] updates related to neo4j --- engine/sessions/session.go | 5 +++++ utils/graph.go | 8 +++++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/engine/sessions/session.go b/engine/sessions/session.go index b85202b8..a6e2240f 100644 --- a/engine/sessions/session.go +++ b/engine/sessions/session.go @@ -23,6 +23,7 @@ import ( assetdb "github.com/owasp-amass/asset-db" "github.com/owasp-amass/asset-db/cache" "github.com/owasp-amass/asset-db/repository" + "github.com/owasp-amass/asset-db/repository/neo4j" "github.com/owasp-amass/asset-db/repository/sqlrepo" ) @@ -168,6 +169,10 @@ func (s *Session) selectDBMS() error { path := filepath.Join(config.OutputDirectory(s.cfg.Dir), "amass.sqlite") s.dsn = path s.dbtype = sqlrepo.SQLite + } else if db.System == "neo4j" || db.System == "neo4+s" || db.System == "neo4j+sec" || + db.System == "bolt" || db.System == "bolt+s" || db.System == "bolt+sec" { + s.dsn = db.URL + s.dbtype = neo4j.Neo4j } // Break the loop once the primary database is found. break diff --git a/utils/graph.go b/utils/graph.go index 578ae37e..9d32b5e8 100644 --- a/utils/graph.go +++ b/utils/graph.go @@ -12,6 +12,7 @@ import ( assetdb "github.com/owasp-amass/asset-db" "github.com/owasp-amass/asset-db/repository" "github.com/owasp-amass/asset-db/repository/sqlrepo" + "github.com/owasp-amass/asset-db/repository/neo4j" ) func OpenGraphDatabase(cfg *config.Config) repository.Repository { @@ -24,9 +25,11 @@ func OpenGraphDatabase(cfg *config.Config) repository.Repository { if db.System == "local" { dbase = NewGraph(db.System, filepath.Join(config.OutputDirectory(cfg.Dir), "amass.sqlite"), db.Options) - } else { + if db.System == "postgres" { connStr := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s", db.Host, db.Port, db.Username, db.Password, db.DBName) dbase = NewGraph(db.System, connStr, db.Options) + } else { + dbase = NewGraph(db.System, db.URL, db.Options) } if dbase != nil { @@ -52,6 +55,9 @@ func NewGraph(system, path string, options string) repository.Repository { case "postgres": dbtype = sqlrepo.Postgres dsn = path + case "bolt": + dbtype = neo4j.Neo4j + dsn = path default: return nil }