Skip to content

Commit

Permalink
updates for asset db v0.16.0
Browse files Browse the repository at this point in the history
  • Loading branch information
@caffix
caffix committed Dec 9, 2024
1 parent e6bbe12 commit 2ea3450
Show file tree
Hide file tree
Showing 10 changed files with 76 additions and 119 deletions.
24 changes: 12 additions & 12 deletions cmd/oam_assoc/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ import (
"github.com/owasp-amass/amass/v4/config"
"github.com/owasp-amass/amass/v4/utils"
"github.com/owasp-amass/amass/v4/utils/afmt"
assetdb "github.com/owasp-amass/asset-db"
"github.com/owasp-amass/asset-db/repository"
dbt "github.com/owasp-amass/asset-db/types"
"github.com/owasp-amass/open-asset-model/domain"
oamreg "github.com/owasp-amass/open-asset-model/registration"
Expand Down Expand Up @@ -191,11 +191,11 @@ func main() {
}
}

func printContactInfo(assoc *dbt.Entity, regrel string, since time.Time, db *assetdb.AssetDB) {
func printContactInfo(assoc *dbt.Entity, regrel string, since time.Time, db repository.Repository) {
var contact *dbt.Entity

if edges, err := db.Repo.OutgoingEdges(assoc, since, regrel); err == nil && len(edges) > 0 {
if a, err := db.Repo.FindEntityById(edges[0].ToEntity.ID); err == nil && a != nil {
if edges, err := db.OutgoingEdges(assoc, since, regrel); err == nil && len(edges) > 0 {
if a, err := db.FindEntityById(edges[0].ToEntity.ID); err == nil && a != nil {
contact = a
}
}
Expand All @@ -204,9 +204,9 @@ func printContactInfo(assoc *dbt.Entity, regrel string, since time.Time, db *ass
}

for _, out := range []string{"person", "organization", "location", "phone", "email"} {
if edges, err := db.Repo.OutgoingEdges(contact, since, out); err == nil && len(edges) > 0 {
if edges, err := db.OutgoingEdges(contact, since, out); err == nil && len(edges) > 0 {
for _, edge := range edges {
if a, err := db.Repo.FindEntityById(edge.ToEntity.ID); err == nil && a != nil {
if a, err := db.FindEntityById(edge.ToEntity.ID); err == nil && a != nil {
fmt.Fprintf(color.Output, "%s%s%s\n",
afmt.Blue(string(a.Asset.AssetType())), afmt.Blue(": "), afmt.Green(a.Asset.Key()))
}
Expand All @@ -215,19 +215,19 @@ func printContactInfo(assoc *dbt.Entity, regrel string, since time.Time, db *ass
}
}

func getAssociations(name string, since time.Time, db *assetdb.AssetDB) []*dbt.Entity {
func getAssociations(name string, since time.Time, db repository.Repository) []*dbt.Entity {
var results []*dbt.Entity

fqdns, err := db.Repo.FindEntityByContent(&domain.FQDN{Name: name}, since)
fqdns, err := db.FindEntityByContent(&domain.FQDN{Name: name}, since)
if err != nil || len(fqdns) == 0 {
return results
}

var assets []*dbt.Entity
for _, fqdn := range fqdns {
if edges, err := db.Repo.OutgoingEdges(fqdn, since, "registration"); err == nil && len(edges) > 0 {
if edges, err := db.OutgoingEdges(fqdn, since, "registration"); err == nil && len(edges) > 0 {
for _, edge := range edges {
if a, err := db.Repo.FindEntityById(edge.ToEntity.ID); err == nil && a != nil {
if a, err := db.FindEntityById(edge.ToEntity.ID); err == nil && a != nil {
assets = append(assets, a)
}
}
Expand All @@ -246,9 +246,9 @@ func getAssociations(name string, since time.Time, db *assetdb.AssetDB) []*dbt.E
findings = []*dbt.Entity{}

for _, a := range assets {
if edges, err := db.Repo.OutgoingEdges(a, since, "associated_with"); err == nil && len(edges) > 0 {
if edges, err := db.OutgoingEdges(a, since, "associated_with"); err == nil && len(edges) > 0 {
for _, edge := range edges {
asset, err := db.Repo.FindEntityById(edge.ToEntity.ID)
asset, err := db.FindEntityById(edge.ToEntity.ID)
if err != nil || asset == nil {
continue
}
Expand Down
12 changes: 6 additions & 6 deletions cmd/oam_subs/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ import (
"github.com/owasp-amass/amass/v4/config"
"github.com/owasp-amass/amass/v4/utils"
"github.com/owasp-amass/amass/v4/utils/afmt"
assetdb "github.com/owasp-amass/asset-db"
"github.com/owasp-amass/asset-db/repository"
dbt "github.com/owasp-amass/asset-db/types"
"github.com/owasp-amass/open-asset-model/domain"
)
Expand Down Expand Up @@ -175,7 +175,7 @@ func main() {
showData(&args, asninfo, db)
}

func showData(args *dbArgs, asninfo bool, db *assetdb.AssetDB) {
func showData(args *dbArgs, asninfo bool, db repository.Repository) {
var total int
var err error
var outfile *os.File
Expand Down Expand Up @@ -265,7 +265,7 @@ func showData(args *dbArgs, asninfo bool, db *assetdb.AssetDB) {
}
}

func getNames(ctx context.Context, domains []string, asninfo bool, db *assetdb.AssetDB) []*utils.Output {
func getNames(ctx context.Context, domains []string, asninfo bool, db repository.Repository) []*utils.Output {
if len(domains) == 0 {
return nil
}
Expand All @@ -276,8 +276,8 @@ func getNames(ctx context.Context, domains []string, asninfo bool, db *assetdb.A

var assets []*dbt.Entity
for _, d := range domains {
if ents, err := db.Repo.FindEntityByContent(&domain.FQDN{Name: d}, qtime); err == nil && len(ents) == 1 {
if n, err := utils.FindByFQDNScope(db.Repo, ents[0], qtime); err == nil && len(n) > 0 {
if ents, err := db.FindEntityByContent(&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...)
}
}
Expand All @@ -296,7 +296,7 @@ func getNames(ctx context.Context, domains []string, asninfo bool, db *assetdb.A
return names
}

func addAddresses(ctx context.Context, db *assetdb.AssetDB, names []*utils.Output, asninfo bool, cache *utils.ASNCache) []*utils.Output {
func addAddresses(ctx context.Context, db repository.Repository, names []*utils.Output, asninfo bool, cache *utils.ASNCache) []*utils.Output {
var namestrs []string
lookup := make(outLookup, len(names))
for _, n := range names {
Expand Down
8 changes: 4 additions & 4 deletions cmd/oam_track/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ import (
"github.com/owasp-amass/amass/v4/config"
"github.com/owasp-amass/amass/v4/utils"
"github.com/owasp-amass/amass/v4/utils/afmt"
assetdb "github.com/owasp-amass/asset-db"
"github.com/owasp-amass/asset-db/repository"
dbt "github.com/owasp-amass/asset-db/types"
"github.com/owasp-amass/open-asset-model/domain"
)
Expand Down Expand Up @@ -153,15 +153,15 @@ func main() {
}
}

func getNewNames(domains []string, since time.Time, db *assetdb.AssetDB) []string {
func getNewNames(domains []string, since time.Time, db repository.Repository) []string {
if len(domains) == 0 {
return []string{}
}

var assets []*dbt.Entity
for _, d := range domains {
if ents, err := db.Repo.FindEntityByContent(&domain.FQDN{Name: d}, since); err == nil && len(ents) == 1 {
if n, err := utils.FindByFQDNScope(db.Repo, ents[0], since); err == nil && len(n) > 0 {
if ents, err := db.FindEntityByContent(&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...)
}
}
Expand Down
4 changes: 2 additions & 2 deletions engine/plugins/service_discovery/http_probes/plugin.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ import (
oam "github.com/owasp-amass/open-asset-model"
oamcert "github.com/owasp-amass/open-asset-model/certificate"
"github.com/owasp-amass/open-asset-model/relation"
"github.com/owasp-amass/open-asset-model/service"
oamserv "github.com/owasp-amass/open-asset-model/service"

Check failure on line 22 in engine/plugins/service_discovery/http_probes/plugin.go

View workflow job for this annotation

GitHub Actions / Lint (ubuntu-latest, 1.23.1) 

"github.com/owasp-amass/open-asset-model/service" imported as oamserv and not used (typecheck)
)

type httpProbing struct {
Expand Down Expand Up @@ -174,7 +174,7 @@ func (hp *httpProbing) store(e *et.Event, resp *http.Response, entity *dbt.Entit
return findings
}

serv = s.Asset.(*service.Service)
serv = s.Asset.(*oamserv.Service)
// for adding the source information
findings = append(findings, &support.Finding{
From: entity,
Expand Down
19 changes: 10 additions & 9 deletions engine/sessions/session.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ type Session struct {
ps *pubsub.Logger
cfg *config.Config
scope *scope.Scope
db *assetdb.AssetDB
db repository.Repository
dsn string
dbtype string
c *cache.Cache
Expand Down Expand Up @@ -71,7 +71,7 @@ func CreateSession(cfg *config.Config) (et.Session, error) {
}
s.tmpdir = dir

s.c, err = cache.New(c, s.db.Repo, time.Minute)
s.c, err = cache.New(c, s.db, time.Minute)
if err != nil || s.c == nil {
return nil, errors.New("failed to create the session cache")
}
Expand All @@ -98,7 +98,7 @@ func (s *Session) Scope() *scope.Scope {
return s.scope
}

func (s *Session) DB() *assetdb.AssetDB {
func (s *Session) DB() repository.Repository {
return s.db
}

Expand Down Expand Up @@ -178,8 +178,8 @@ func (s *Session) selectDBMS() error {
return errors.New("no primary database specified in the configuration")
}
// Initialize the database store
store := assetdb.New(s.dbtype, s.dsn)
if store == nil {
store, err := assetdb.New(s.dbtype, s.dsn)
if err != nil {
return errors.New("failed to initialize database store")
}
s.db = store
Expand All @@ -192,10 +192,11 @@ func createFileCacheRepo() (repository.Repository, string, error) {
return nil, "", errors.New("failed to create the temp dir")
}

c := assetdb.New(sqlrepo.SQLite, filepath.Join(dir, "cache.sqlite"))
if c == nil {
return nil, "", errors.New("failed to create the cache db")
//c := assetdb.New(sqlrepo.SQLite, filepath.Join(dir, "cache.sqlite"))
c, err := assetdb.New(sqlrepo.SQLiteMemory, "")
if err != nil {
return nil, "", fmt.Errorf("failed to create the cache db: %s", err.Error())
}

return c.Repo, dir, nil
return c, dir, nil
}
4 changes: 2 additions & 2 deletions engine/types/sessions.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,8 @@ import (
"github.com/owasp-amass/amass/v4/config"
"github.com/owasp-amass/amass/v4/engine/pubsub"
"github.com/owasp-amass/amass/v4/engine/sessions/scope"
assetdb "github.com/owasp-amass/asset-db"
"github.com/owasp-amass/asset-db/cache"
"github.com/owasp-amass/asset-db/repository"
)

type Session interface {
Expand All @@ -23,7 +23,7 @@ type Session interface {
PubSub() *pubsub.Logger
Config() *config.Config
Scope() *scope.Scope
DB() *assetdb.AssetDB
DB() repository.Repository
Cache() *cache.Cache
TmpDir() string
Stats() *SessionStats
Expand Down
23 changes: 11 additions & 12 deletions utils/addr.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,26 +9,25 @@ import (
"time"

"github.com/caffix/stringset"
assetdb "github.com/owasp-amass/asset-db"
"github.com/owasp-amass/asset-db/repository"
dbt "github.com/owasp-amass/asset-db/types"
"github.com/owasp-amass/open-asset-model/domain"
"github.com/owasp-amass/open-asset-model/network"
"github.com/owasp-amass/open-asset-model/relation"
)

func ReadASPrefixes(db *assetdb.AssetDB, asn int, since time.Time) []string {
func ReadASPrefixes(db repository.Repository, asn int, since time.Time) []string {
var prefixes []string

fqdns, err := db.FindByContent(&network.AutonomousSystem{Number: asn}, since)
fqdns, err := db.FindEntityByContent(&network.AutonomousSystem{Number: asn}, since)
if err != nil || len(fqdns) != 1 {
return prefixes
}
fqdn := fqdns[0]

if edges, err := db.Repo.OutgoingEdges(fqdn, since, "announces"); err == nil && len(edges) > 0 {
if edges, err := db.OutgoingEdges(fqdn, since, "announces"); err == nil && len(edges) > 0 {
for _, edge := range edges {
if a, err := db.Repo.FindEntityById(edge.ToEntity.ID); err != nil {
if a, err := db.FindEntityById(edge.ToEntity.ID); err != nil {
continue
} else if netblock, ok := a.Asset.(*network.Netblock); ok {
prefixes = append(prefixes, netblock.CIDR.String())
Expand All @@ -44,10 +43,10 @@ type NameAddrPair struct {
Addr *network.IPAddress
}

func NamesToAddrs(db *assetdb.AssetDB, since time.Time, names ...string) ([]*NameAddrPair, error) {
func NamesToAddrs(db repository.Repository, since time.Time, names ...string) ([]*NameAddrPair, error) {
var fqdns []*dbt.Entity
for _, name := range names {
if ents, err := db.Repo.FindEntityByContent(&domain.FQDN{Name: name}, since); err == nil && len(ents) == 1 {
if ents, err := db.FindEntityByContent(&domain.FQDN{Name: name}, since); err == nil && len(ents) == 1 {
fqdns = append(fqdns, ents[0])
}
}
Expand All @@ -56,20 +55,20 @@ func NamesToAddrs(db *assetdb.AssetDB, since time.Time, names ...string) ([]*Nam
// get the IPs associated with SRV, NS, and MX records
loop:
for _, fqdn := range fqdns {
if edges, err := db.Repo.OutgoingEdges(fqdn, since, "dns_record"); err == nil && len(edges) > 0 {
if edges, err := db.OutgoingEdges(fqdn, since, "dns_record"); err == nil && len(edges) > 0 {
for _, edge := range edges {
switch v := edge.Relation.(type) {
case *relation.BasicDNSRelation:
if v.Header.RRType == 1 || v.Header.RRType == 28 {
if ip, err := getAddr(db.Repo, edge.ToEntity, since); err == nil {
if ip, err := getAddr(db, edge.ToEntity, since); err == nil {
results = append(results, &NameAddrPair{
FQDN: fqdn.Asset.(*domain.FQDN),
Addr: ip,
})
continue loop
}
} else if v.Header.RRType == 5 {
if ip, err := cnameQuery(db.Repo, edge.ToEntity, since); err == nil {
if ip, err := cnameQuery(db, edge.ToEntity, since); err == nil {
results = append(results, &NameAddrPair{
FQDN: fqdn.Asset.(*domain.FQDN),
Addr: ip,
Expand All @@ -79,7 +78,7 @@ loop:
}
case *relation.PrefDNSRelation:
if v.Header.RRType == 2 || v.Header.RRType == 15 {
if ip, err := oneMoreName(db.Repo, edge.ToEntity, since); err == nil {
if ip, err := oneMoreName(db, edge.ToEntity, since); err == nil {
results = append(results, &NameAddrPair{
FQDN: fqdn.Asset.(*domain.FQDN),
Addr: ip,
Expand All @@ -89,7 +88,7 @@ loop:
}
case *relation.SRVDNSRelation:
if v.Header.RRType == 33 {
if ip, err := oneMoreName(db.Repo, edge.ToEntity, since); err == nil {
if ip, err := oneMoreName(db, edge.ToEntity, since); err == nil {
results = append(results, &NameAddrPair{
FQDN: fqdn.Asset.(*domain.FQDN),
Addr: ip,
Expand Down
10 changes: 5 additions & 5 deletions utils/asncache.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ import (
"time"

"github.com/caffix/stringset"
assetdb "github.com/owasp-amass/asset-db"
"github.com/owasp-amass/asset-db/repository"
oam "github.com/owasp-amass/open-asset-model"
"github.com/owasp-amass/open-asset-model/network"
oamreg "github.com/owasp-amass/open-asset-model/registration"
Expand Down Expand Up @@ -222,9 +222,9 @@ func (c *ASNCache) AddrSearch(addr string) *ASNRequest {
}
}

func FillCache(cache *ASNCache, db *assetdb.AssetDB) error {
func FillCache(cache *ASNCache, db repository.Repository) error {
start := time.Now().Add(-730 * time.Hour)
assets, err := db.Repo.FindEntitiesByType(oam.AutonomousSystem, start)
assets, err := db.FindEntitiesByType(oam.AutonomousSystem, start)
if err != nil {
return err
}
Expand All @@ -236,13 +236,13 @@ func FillCache(cache *ASNCache, db *assetdb.AssetDB) error {
}

var desc string
edges, err := db.Repo.OutgoingEdges(a, start, "registration")
edges, err := db.OutgoingEdges(a, start, "registration")
if err != nil || len(edges) == 0 {
continue
}

for _, edge := range edges {
if asset, err := db.Repo.FindEntityById(edge.ToEntity.ID); err == nil && asset != nil {
if asset, err := db.FindEntityById(edge.ToEntity.ID); err == nil && asset != nil {
if autnum, ok := asset.Asset.(*oamreg.AutnumRecord); ok && autnum != nil {
desc = autnum.Handle + " - " + autnum.Name
break
Expand Down
Loading

0 comments on commit 2ea3450

Please sign in to comment.