updates related to neo4j

Author: caffix

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

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
 	}