updates related to the changes to assetdb

caffix · caffix · commit 35bad0a16b7a · 2024-12-21T13:37:56.000-05:00
diff --git a/cmd/oam_assoc/main.go b/cmd/oam_assoc/main.go
@@ -218,7 +218,7 @@ func printContactInfo(assoc *dbt.Entity, regrel string, since time.Time, db repo
 func getAssociations(name string, since time.Time, db repository.Repository) []*dbt.Entity {
 	var results []*dbt.Entity
 
-	fqdns, err := db.FindEntityByContent(&domain.FQDN{Name: name}, since)
+	fqdns, err := db.FindEntitiesByContent(&domain.FQDN{Name: name}, since)
 	if err != nil || len(fqdns) == 0 {
 		return results
 	}
diff --git a/cmd/oam_subs/main.go b/cmd/oam_subs/main.go
@@ -276,7 +276,7 @@ func getNames(ctx context.Context, domains []string, asninfo bool, db repository
 
 	var assets []*dbt.Entity
 	for _, d := range domains {
-		if ents, err := db.FindEntityByContent(&domain.FQDN{Name: d}, qtime); err == nil && len(ents) == 1 {
+		if ents, err := db.FindEntitiesByContent(&domain.FQDN{Name: d}, qtime); err == nil && len(ents) == 1 {
 			if n, err := utils.FindByFQDNScope(db, ents[0], qtime); err == nil && len(n) > 0 {
 				assets = append(assets, n...)
 			}
diff --git a/cmd/oam_track/main.go b/cmd/oam_track/main.go
@@ -160,7 +160,7 @@ func getNewNames(domains []string, since time.Time, db repository.Repository) []
 
 	var assets []*dbt.Entity
 	for _, d := range domains {
-		if ents, err := db.FindEntityByContent(&domain.FQDN{Name: d}, since); err == nil && len(ents) == 1 {
+		if ents, err := db.FindEntitiesByContent(&domain.FQDN{Name: d}, since); err == nil && len(ents) == 1 {
 			if n, err := utils.FindByFQDNScope(db, ents[0], since); err == nil && len(n) > 0 {
 				assets = append(assets, n...)
 			}
diff --git a/engine/plugins/dns/apex.go b/engine/plugins/dns/apex.go
@@ -37,7 +37,7 @@ func (d *dnsApex) check(e *et.Event) error {
 	for _, name := range d.plugin.apexList.Slice() {
 		if idx := strings.Index(fqdn.Name, name); idx != -1 && idx != 0 && idx < best {
 			best = idx
-			if ents, err := e.Session.Cache().FindEntityByContent(
+			if ents, err := e.Session.Cache().FindEntitiesByContent(
 				&domain.FQDN{Name: name}, e.Session.Cache().StartTime()); err == nil && len(ents) == 1 {
 				apex = ents[0]
 			}
diff --git a/engine/plugins/dns/plugin.go b/engine/plugins/dns/plugin.go
@@ -142,7 +142,7 @@ func (d *dnsPlugin) lookupWithinTTL(session et.Session, name string, atype oam.A
 		return results
 	}
 
-	ents, err := session.Cache().FindEntityByContent(&domain.FQDN{Name: name}, time.Time{})
+	ents, err := session.Cache().FindEntitiesByContent(&domain.FQDN{Name: name}, time.Time{})
 	if err != nil || len(ents) != 1 {
 		return results
 	}
diff --git a/engine/plugins/dns/subs.go b/engine/plugins/dns/subs.go
@@ -98,7 +98,7 @@ func (d *dnsSubs) registered(e *et.Event, name string) string {
 		}
 	}
 
-	fqdns, err := e.Session.Cache().FindEntityByContent(&domain.FQDN{Name: name}, time.Time{})
+	fqdns, err := e.Session.Cache().FindEntitiesByContent(&domain.FQDN{Name: name}, time.Time{})
 	if err != nil || len(fqdns) != 1 {
 		return ""
 	}
@@ -166,7 +166,7 @@ func (d *dnsSubs) traverse(e *et.Event, dom string, fqdn *dbt.Entity, src *et.So
 func (d *dnsSubs) lookup(e *et.Event, subdomain string, since time.Time) []*relSubs {
 	var alias []*relSubs
 
-	fqdns, err := e.Session.Cache().FindEntityByContent(&domain.FQDN{Name: subdomain}, time.Time{})
+	fqdns, err := e.Session.Cache().FindEntitiesByContent(&domain.FQDN{Name: subdomain}, time.Time{})
 	if err != nil || len(fqdns) != 1 {
 		return alias
 	}
diff --git a/engine/plugins/horizontals/fqdn.go b/engine/plugins/horizontals/fqdn.go
@@ -77,7 +77,7 @@ func (h *horfqdn) check(e *et.Event) error {
 
 		var assets []*dbt.Entity
 		for _, im := range impacted {
-			if a, err := e.Session.Cache().FindEntityByContent(im.Asset,
+			if a, err := e.Session.Cache().FindEntitiesByContent(im.Asset,
 				e.Session.Cache().StartTime()); err == nil && len(a) == 1 {
 				assets = append(assets, a[0])
 			} else if n := h.store(e, im.Asset); n != nil {
diff --git a/engine/plugins/horizontals/plugin.go b/engine/plugins/horizontals/plugin.go
@@ -98,7 +98,7 @@ func (h *horizPlugin) addAssociatedRelationship(e *et.Event, assocs []*scope.Ass
 			}
 
 			if match, result := e.Session.Scope().IsAssetInScope(impacted.Asset, conf); result >= conf && match != nil {
-				if a, err := e.Session.Cache().FindEntityByContent(match, e.Session.Cache().StartTime()); err == nil && len(a) == 1 {
+				if a, err := e.Session.Cache().FindEntitiesByContent(match, e.Session.Cache().StartTime()); err == nil && len(a) == 1 {
 					for _, assoc2 := range e.Session.Scope().AssetsWithAssociation(e.Session.Cache(), a[0]) {
 						h.makeAssocRelationshipEntries(e, assoc.Match, assoc2)
 					}
@@ -135,7 +135,7 @@ func (h *horizPlugin) process(e *et.Event, assets []*dbt.Entity) {
 			h.sweepAroundIPs(e, asset)
 			//h.sweepNetblock(e, v, src)
 		case *oamreg.IPNetRecord:
-			if ents, err := e.Session.Cache().FindEntityByContent(
+			if ents, err := e.Session.Cache().FindEntitiesByContent(
 				&oamnet.Netblock{CIDR: v.CIDR, Type: v.Type}, e.Session.Cache().StartTime()); err == nil && len(ents) == 1 {
 				a := ents[0]
 
@@ -173,7 +173,7 @@ func (h *horizPlugin) ipPTRTargetsInScope(e *et.Event, nb *dbt.Entity) {
 				continue
 			}
 
-			if ents, err := e.Session.Cache().FindEntityByContent(
+			if ents, err := e.Session.Cache().FindEntitiesByContent(
 				&domain.FQDN{Name: resolve.RemoveLastDot(reverse)}, e.Session.Cache().StartTime()); err == nil && len(ents) == 1 {
 				a := ents[0]
 
diff --git a/engine/plugins/support/database.go b/engine/plugins/support/database.go
@@ -27,31 +27,31 @@ func SourceToAssetsWithinTTL(session et.Session, name, atype string, src *et.Sou
 
 	switch atype {
 	case string(oam.FQDN):
-		roots, err := session.Cache().FindEntityByContent(&domain.FQDN{Name: name}, since)
+		roots, err := session.Cache().FindEntitiesByContent(&domain.FQDN{Name: name}, since)
 		if err != nil || len(roots) != 1 {
 			return nil
 		}
 		root := roots[0]
 
 		entities, _ = utils.FindByFQDNScope(session.Cache(), root, since)
 	case string(oam.EmailAddress):
-		entities, _ = session.Cache().FindEntityByContent(EmailToOAMEmailAddress(name), since)
+		entities, _ = session.Cache().FindEntitiesByContent(EmailToOAMEmailAddress(name), since)
 	case string(oam.AutnumRecord):
 		num, err := strconv.Atoi(name)
 		if err != nil {
 			return nil
 		}
 
-		entities, _ = session.Cache().FindEntityByContent(&oamreg.AutnumRecord{Number: num}, since)
+		entities, _ = session.Cache().FindEntitiesByContent(&oamreg.AutnumRecord{Number: num}, since)
 	case string(oam.IPNetRecord):
 		prefix, err := netip.ParsePrefix(name)
 		if err != nil {
 			return nil
 		}
 
-		entities, _ = session.Cache().FindEntityByContent(&oamreg.IPNetRecord{CIDR: prefix}, since)
+		entities, _ = session.Cache().FindEntitiesByContent(&oamreg.IPNetRecord{CIDR: prefix}, since)
 	case string(oam.Service):
-		entities, _ = session.Cache().FindEntityByContent(&service.Service{Identifier: name}, since)
+		entities, _ = session.Cache().FindEntitiesByContent(&service.Service{Identifier: name}, since)
 	}
 
 	var results []*dbt.Entity
diff --git a/engine/plugins/support/support.go b/engine/plugins/support/support.go
@@ -249,7 +249,7 @@ func cleanSessionNetblocks() {
 
 func IPAddressSweep(e *et.Event, addr *oamnet.IPAddress, src *et.Source, size int, callback SweepCallback) {
 	// do not work on an IP address that was processed previously
-	_, err := e.Session.Cache().FindEntityByContent(addr, e.Session.Cache().StartTime())
+	_, err := e.Session.Cache().FindEntitiesByContent(addr, e.Session.Cache().StartTime())
 	if err == nil {
 		return
 	}
@@ -287,7 +287,7 @@ func IPAddressSweep(e *et.Event, addr *oamnet.IPAddress, src *et.Source, size in
 }
 
 func IsCNAME(session et.Session, name *domain.FQDN) (*domain.FQDN, bool) {
-	fqdns, err := session.Cache().FindEntityByContent(name, session.Cache().StartTime())
+	fqdns, err := session.Cache().FindEntitiesByContent(name, session.Cache().StartTime())
 	if err != nil || len(fqdns) != 1 {
 		return nil, false
 	}
@@ -308,7 +308,7 @@ func IsCNAME(session et.Session, name *domain.FQDN) (*domain.FQDN, bool) {
 }
 
 func NameIPAddresses(session et.Session, name *domain.FQDN) []*oamnet.IPAddress {
-	fqdns, err := session.Cache().FindEntityByContent(name, session.Cache().StartTime())
+	fqdns, err := session.Cache().FindEntitiesByContent(name, session.Cache().StartTime())
 	if err != nil || len(fqdns) != 1 {
 		return nil
 	}
diff --git a/engine/sessions/scope/assoc.go b/engine/sessions/scope/assoc.go
@@ -365,7 +365,7 @@ func (s *Scope) IsAddressInScope(c *cache.Cache, ip *oamnet.IPAddress) bool {
 		return true
 	}
 
-	addrs, err := c.FindEntityByContent(ip, c.StartTime())
+	addrs, err := c.FindEntitiesByContent(ip, c.StartTime())
 	if err != nil || len(addrs) != 1 {
 		return false
 	}
diff --git a/utils/addr.go b/utils/addr.go
@@ -19,7 +19,7 @@ import (
 func ReadASPrefixes(db repository.Repository, asn int, since time.Time) []string {
 	var prefixes []string
 
-	fqdns, err := db.FindEntityByContent(&network.AutonomousSystem{Number: asn}, since)
+	fqdns, err := db.FindEntitiesByContent(&network.AutonomousSystem{Number: asn}, since)
 	if err != nil || len(fqdns) != 1 {
 		return prefixes
 	}
@@ -46,7 +46,7 @@ type NameAddrPair struct {
 func NamesToAddrs(db repository.Repository, since time.Time, names ...string) ([]*NameAddrPair, error) {
 	var fqdns []*dbt.Entity
 	for _, name := range names {
-		if ents, err := db.FindEntityByContent(&domain.FQDN{Name: name}, since); err == nil && len(ents) == 1 {
+		if ents, err := db.FindEntitiesByContent(&domain.FQDN{Name: name}, since); err == nil && len(ents) == 1 {
 			fqdns = append(fqdns, ents[0])
 		}
 	}
diff --git a/utils/graph.go b/utils/graph.go
@@ -11,8 +11,8 @@ import (
 	"github.com/owasp-amass/amass/v4/config"
 	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"
+	"github.com/owasp-amass/asset-db/repository/sqlrepo"
 )
 
 func OpenGraphDatabase(cfg *config.Config) repository.Repository {
@@ -25,7 +25,7 @@ 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)
-			if db.System == "postgres" {
+			} 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 {
diff --git a/utils/viz/viz.go b/utils/viz/viz.go
@@ -42,7 +42,7 @@ func VizData(domains []string, since time.Time, db repository.Repository) ([]Nod
 
 	var next []*types.Entity
 	for _, d := range domains {
-		if ents, err := db.FindEntityByContent(&domain.FQDN{Name: d}, since); err == nil && len(ents) == 1 {
+		if ents, err := db.FindEntitiesByContent(&domain.FQDN{Name: d}, since); err == nil && len(ents) == 1 {
 			if n, err := utils.FindByFQDNScope(db, ents[0], since); err == nil && len(n) > 0 {
 				next = append(next, n...)
 			}

Original file line number	Diff line number	Diff line change
`@@ -218,7 +218,7 @@ func printContactInfo(assoc *dbt.Entity, regrel string, since time.Time, db repo`
`218`	`218`	`func getAssociations(name string, since time.Time, db repository.Repository) []*dbt.Entity {`
`219`	`219`	`var results []*dbt.Entity`
`220`	`220`
`221`		`- fqdns, err := db.FindEntityByContent(&domain.FQDN{Name: name}, since)`
	`221`	`+ fqdns, err := db.FindEntitiesByContent(&domain.FQDN{Name: name}, since)`
`222`	`222`	`if err != nil \|\| len(fqdns) == 0 {`
`223`	`223`	`return results`
`224`	`224`	`}`
Original file line number	Diff line number	Diff line change
`@@ -276,7 +276,7 @@ func getNames(ctx context.Context, domains []string, asninfo bool, db repository`
`276`	`276`
`277`	`277`	`var assets []*dbt.Entity`
`278`	`278`	`for _, d := range domains {`
`279`		`- if ents, err := db.FindEntityByContent(&domain.FQDN{Name: d}, qtime); err == nil && len(ents) == 1 {`
	`279`	`+ if ents, err := db.FindEntitiesByContent(&domain.FQDN{Name: d}, qtime); err == nil && len(ents) == 1 {`
`280`	`280`	`if n, err := utils.FindByFQDNScope(db, ents[0], qtime); err == nil && len(n) > 0 {`
`281`	`281`	`assets = append(assets, n...)`
`282`	`282`	`}`
Original file line number	Diff line number	Diff line change
`@@ -160,7 +160,7 @@ func getNewNames(domains []string, since time.Time, db repository.Repository) []`
`160`	`160`
`161`	`161`	`var assets []*dbt.Entity`
`162`	`162`	`for _, d := range domains {`
`163`		`- if ents, err := db.FindEntityByContent(&domain.FQDN{Name: d}, since); err == nil && len(ents) == 1 {`
	`163`	`+ if ents, err := db.FindEntitiesByContent(&domain.FQDN{Name: d}, since); err == nil && len(ents) == 1 {`
`164`	`164`	`if n, err := utils.FindByFQDNScope(db, ents[0], since); err == nil && len(n) > 0 {`
`165`	`165`	`assets = append(assets, n...)`
`166`	`166`	`}`
Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,7 @@ func (d dnsApex) check(e et.Event) error {`
`37`	`37`	`for _, name := range d.plugin.apexList.Slice() {`
`38`	`38`	`if idx := strings.Index(fqdn.Name, name); idx != -1 && idx != 0 && idx < best {`
`39`	`39`	`best = idx`
`40`		`- if ents, err := e.Session.Cache().FindEntityByContent(`
	`40`	`+ if ents, err := e.Session.Cache().FindEntitiesByContent(`
`41`	`41`	`&domain.FQDN{Name: name}, e.Session.Cache().StartTime()); err == nil && len(ents) == 1 {`
`42`	`42`	`apex = ents[0]`
`43`	`43`	`}`
Original file line number	Diff line number	Diff line change
`@@ -142,7 +142,7 @@ func (d *dnsPlugin) lookupWithinTTL(session et.Session, name string, atype oam.A`
`142`	`142`	`return results`
`143`	`143`	`}`
`144`	`144`
`145`		`- ents, err := session.Cache().FindEntityByContent(&domain.FQDN{Name: name}, time.Time{})`
	`145`	`+ ents, err := session.Cache().FindEntitiesByContent(&domain.FQDN{Name: name}, time.Time{})`
`146`	`146`	`if err != nil \|\| len(ents) != 1 {`
`147`	`147`	`return results`
`148`	`148`	`}`
Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,7 @@ func (d dnsSubs) registered(e et.Event, name string) string {`
`98`	`98`	`}`
`99`	`99`	`}`
`100`	`100`
`101`		`- fqdns, err := e.Session.Cache().FindEntityByContent(&domain.FQDN{Name: name}, time.Time{})`
	`101`	`+ fqdns, err := e.Session.Cache().FindEntitiesByContent(&domain.FQDN{Name: name}, time.Time{})`
`102`	`102`	`if err != nil \|\| len(fqdns) != 1 {`
`103`	`103`	`return ""`
`104`	`104`	`}`
`@@ -166,7 +166,7 @@ func (d dnsSubs) traverse(e et.Event, dom string, fqdn dbt.Entity, src et.So`
`166`	`166`	`func (d dnsSubs) lookup(e et.Event, subdomain string, since time.Time) []*relSubs {`
`167`	`167`	`var alias []*relSubs`
`168`	`168`
`169`		`- fqdns, err := e.Session.Cache().FindEntityByContent(&domain.FQDN{Name: subdomain}, time.Time{})`
	`169`	`+ fqdns, err := e.Session.Cache().FindEntitiesByContent(&domain.FQDN{Name: subdomain}, time.Time{})`
`170`	`170`	`if err != nil \|\| len(fqdns) != 1 {`
`171`	`171`	`return alias`
`172`	`172`	`}`
Original file line number	Diff line number	Diff line change
`@@ -249,7 +249,7 @@ func cleanSessionNetblocks() {`
`249`	`249`
`250`	`250`	`func IPAddressSweep(e et.Event, addr oamnet.IPAddress, src *et.Source, size int, callback SweepCallback) {`
`251`	`251`	`// do not work on an IP address that was processed previously`
`252`		`- _, err := e.Session.Cache().FindEntityByContent(addr, e.Session.Cache().StartTime())`
	`252`	`+ _, err := e.Session.Cache().FindEntitiesByContent(addr, e.Session.Cache().StartTime())`
`253`	`253`	`if err == nil {`
`254`	`254`	`return`
`255`	`255`	`}`
`@@ -287,7 +287,7 @@ func IPAddressSweep(e et.Event, addr oamnet.IPAddress, src *et.Source, size in`
`287`	`287`	`}`
`288`	`288`
`289`	`289`	`func IsCNAME(session et.Session, name domain.FQDN) (domain.FQDN, bool) {`
`290`		`- fqdns, err := session.Cache().FindEntityByContent(name, session.Cache().StartTime())`
	`290`	`+ fqdns, err := session.Cache().FindEntitiesByContent(name, session.Cache().StartTime())`
`291`	`291`	`if err != nil \|\| len(fqdns) != 1 {`
`292`	`292`	`return nil, false`
`293`	`293`	`}`
`@@ -308,7 +308,7 @@ func IsCNAME(session et.Session, name domain.FQDN) (domain.FQDN, bool) {`
`308`	`308`	`}`
`309`	`309`
`310`	`310`	`func NameIPAddresses(session et.Session, name domain.FQDN) []oamnet.IPAddress {`
`311`		`- fqdns, err := session.Cache().FindEntityByContent(name, session.Cache().StartTime())`
	`311`	`+ fqdns, err := session.Cache().FindEntitiesByContent(name, session.Cache().StartTime())`
`312`	`312`	`if err != nil \|\| len(fqdns) != 1 {`
`313`	`313`	`return nil`
`314`	`314`	`}`