Skip to content

Commit

Permalink
updates related to neo4j
Browse files Browse the repository at this point in the history
  • Loading branch information
caffix committed Dec 21, 2024
1 parent b1c6c15 commit a887622
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 1 deletion.
5 changes: 5 additions & 0 deletions engine/sessions/session.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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
Expand Down
8 changes: 7 additions & 1 deletion utils/graph.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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 {
Expand All @@ -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
}
Expand Down

0 comments on commit a887622

Please sign in to comment.