updated to OAM v0.7.0

Author: caffix

cmd/amass/emails.go

@@ -1,12 +1,11 @@
-// Copyright © by Jeff Foley 2024. All rights reserved.
+// Copyright © by Jeff Foley 2017-2024. All rights reserved.
 // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
 // SPDX-License-Identifier: Apache-2.0
 
 package main
 
 import (
 	"bytes"
-	"context"
 	"flag"
 	"io"
 	"os"
@@ -15,8 +14,8 @@ import (
 
 	"github.com/caffix/stringset"
 	"github.com/fatih/color"
+	assetdb "github.com/owasp-amass/asset-db"
 	"github.com/owasp-amass/config/config"
-	"github.com/owasp-amass/engine/graph"
 	oam "github.com/owasp-amass/open-asset-model"
 	"github.com/owasp-amass/open-asset-model/contact"
 	"github.com/owasp-amass/open-asset-model/domain"
@@ -119,7 +118,7 @@ func runEmailsCommand(clArgs []string) {
 	showEmails(&args, db)
 }
 
-func showEmails(args *emailsArgs, db *graph.Graph) {
+func showEmails(args *emailsArgs, db *assetdb.AssetDB) {
 	var err error
 	var outfile *os.File
 	domains := args.Domains.Slice()
@@ -138,7 +137,7 @@ func showEmails(args *emailsArgs, db *graph.Graph) {
 		_, _ = outfile.Seek(0, 0)
 	}
 
-	addrs := getAddresses(context.Background(), domains, db)
+	addrs := getAddresses(db, domains)
 	if len(addrs) == 0 {
 		r.Println("No email addresses were discovered")
 		return
@@ -149,7 +148,7 @@ func showEmails(args *emailsArgs, db *graph.Graph) {
 	}
 }
 
-func getAddresses(ctx context.Context, domains []string, g *graph.Graph) []string {
+func getAddresses(db *assetdb.AssetDB, domains []string) []string {
 	if len(domains) == 0 {
 		return nil
 	}
@@ -163,7 +162,7 @@ func getAddresses(ctx context.Context, domains []string, g *graph.Graph) []strin
 		fqdns = append(fqdns, &domain.FQDN{Name: d})
 	}
 
-	assets, err := g.DB.FindByScope(fqdns, qtime)
+	assets, err := db.FindByScope(fqdns, qtime)
 	if err != nil {
 		return nil
 	}

cmd/amass/main.go

@@ -25,10 +25,12 @@ package main
 import (
 	"bufio"
 	"bytes"
+	"embed"
 	"flag"
 	"fmt"
 	"io"
 	"log/slog"
+	"math/rand"
 	"net"
 	"net/netip"
 	"os"
@@ -39,14 +41,22 @@ import (
 
 	"github.com/caffix/stringset"
 	"github.com/fatih/color"
+	"github.com/glebarez/sqlite"
 	"github.com/owasp-amass/amass/v4/format"
+	assetdb "github.com/owasp-amass/asset-db"
+	db "github.com/owasp-amass/asset-db"
+	pgmigrations "github.com/owasp-amass/asset-db/migrations/postgres"
+	sqlitemigrations "github.com/owasp-amass/asset-db/migrations/sqlite3"
+	"github.com/owasp-amass/asset-db/repository"
 	"github.com/owasp-amass/config/config"
-	"github.com/owasp-amass/engine/graph"
 	et "github.com/owasp-amass/engine/types"
 	"github.com/owasp-amass/open-asset-model/domain"
 	oamnet "github.com/owasp-amass/open-asset-model/network"
+	migrate "github.com/rubenv/sql-migrate"
 	slogcommon "github.com/samber/slog-common"
 	slogsyslog "github.com/samber/slog-syslog/v2"
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
 )
 
 const (
@@ -153,29 +163,88 @@ func createOutputDirectory(cfg *config.Config) {
 	}
 }
 
-func openGraphDatabase(cfg *config.Config) *graph.Graph {
+func openGraphDatabase(cfg *config.Config) *assetdb.AssetDB {
 	// Add the local database settings to the configuration
 	cfg.GraphDBs = append(cfg.GraphDBs, cfg.LocalDatabaseSettings(cfg.GraphDBs))
 
 	for _, db := range cfg.GraphDBs {
 		if db.Primary {
-			var g *graph.Graph
+			var dbase *assetdb.AssetDB
 
 			if db.System == "local" {
-				g = graph.NewGraph(db.System, filepath.Join(config.OutputDirectory(cfg.Dir), "amass.sqlite"), db.Options)
+				dbase = NewGraph(db.System, filepath.Join(config.OutputDirectory(cfg.Dir), "amass.sqlite"), db.Options)
 			} else {
 				connStr := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s", db.Host, db.Port, db.Username, db.Password, db.DBName)
-				g = graph.NewGraph(db.System, connStr, db.Options)
+				dbase = NewGraph(db.System, connStr, db.Options)
 			}
 
-			if g != nil {
-				return g
+			if dbase != nil {
+				return dbase
 			}
 			break
 		}
 	}
 
-	return graph.NewGraph("memory", "", "")
+	return NewGraph("memory", "", "")
+}
+
+func NewGraph(system, path string, options string) *assetdb.AssetDB {
+	var dsn string
+	var dbtype repository.DBType
+
+	switch system {
+	case "memory":
+		dbtype = repository.SQLite
+		dsn = fmt.Sprintf("file:sqlite%d?mode=memory&cache=shared", rand.Int31n(100))
+	case "local":
+		dbtype = repository.SQLite
+		dsn = path
+	case "postgres":
+		dbtype = repository.Postgres
+		dsn = path
+	default:
+		return nil
+	}
+
+	store := db.New(dbtype, dsn)
+	if store == nil {
+		return nil
+	}
+
+	var name string
+	var fs embed.FS
+	var database gorm.Dialector
+	switch dbtype {
+	case repository.SQLite:
+		name = "sqlite3"
+		fs = sqlitemigrations.Migrations()
+		database = sqlite.Open(dsn)
+	case repository.Postgres:
+		name = "postgres"
+		fs = pgmigrations.Migrations()
+		database = postgres.Open(dsn)
+	}
+
+	sql, err := gorm.Open(database, &gorm.Config{})
+	if err != nil {
+		return nil
+	}
+
+	migrationsSource := migrate.EmbedFileSystemMigrationSource{
+		FileSystem: fs,
+		Root:       "/",
+	}
+
+	sqlDb, err := sql.DB()
+	if err != nil {
+		panic(err)
+	}
+
+	_, err = migrate.Exec(sqlDb, name, migrationsSource, migrate.Up)
+	if err != nil {
+		panic(err)
+	}
+	return store
 }
 
 func getWordList(reader io.Reader) ([]string, error) {