diff --git a/assetdb.go b/assetdb.go index fc7e0c9..8db4f82 100644 --- a/assetdb.go +++ b/assetdb.go @@ -62,14 +62,6 @@ func (as *AssetDB) FindEntityById(id string) (*types.Entity, error) { return as.Repo.FindEntityById(id) } -// FindByScope finds entities in the database by applying all the scope constraints provided -// and last seen after the since parameter. -// If since.IsZero(), the parameter will be ignored. -// It returns the matching entities and an error, if any. -func (as *AssetDB) FindByScope(constraints []oam.Asset, since time.Time) ([]*types.Entity, error) { - return as.Repo.FindEntitiesByScope(constraints, since) -} - // FindEntitiesByType finds all entities in the database of the provided asset type and last seen after the since parameter. // If since.IsZero(), the parameter will be ignored. // It returns the matching entities and an error, if any. diff --git a/assetdb_test.go b/assetdb_test.go index 159a6d3..fed10dd 100644 --- a/assetdb_test.go +++ b/assetdb_test.go @@ -150,36 +150,6 @@ func TestAssetDB(t *testing.T) { } }) - t.Run("FindByScope", func(t *testing.T) { - testCases := []struct { - description string - assets []oam.Asset - expected []*types.Entity - expectedError error - }{ - {"an entity is found", []oam.Asset{&domain.FQDN{Name: "domain.com"}}, []*types.Entity{{ID: "1", Asset: &domain.FQDN{Name: "www.domain.com"}}}, nil}, - {"an entity is not found", []oam.Asset{&domain.FQDN{Name: "domain.com"}}, []*types.Entity{}, fmt.Errorf("entity not found")}, - } - - for _, tc := range testCases { - t.Run(tc.description, func(t *testing.T) { - mockAssetDB := new(mockAssetDB) - adb := AssetDB{ - Repo: mockAssetDB, - } - - mockAssetDB.On("FindEntitiesByScope", tc.assets, start).Return(tc.expected, tc.expectedError) - - result, err := adb.FindByScope(tc.assets, start) - - assert.Equal(t, tc.expected, result) - assert.Equal(t, tc.expectedError, err) - - mockAssetDB.AssertExpectations(t) - }) - } - }) - t.Run("FindEntitiesByType", func(t *testing.T) { testCases := []struct { description string @@ -444,11 +414,6 @@ func (m *mockAssetDB) FindEntityByContent(asset oam.Asset, since time.Time) ([]* return args.Get(0).([]*types.Entity), args.Error(1) } -func (m *mockAssetDB) FindEntitiesByScope(constraints []oam.Asset, since time.Time) ([]*types.Entity, error) { - args := m.Called(constraints, since) - return args.Get(0).([]*types.Entity), args.Error(1) -} - func (m *mockAssetDB) FindEntitiesByType(atype oam.AssetType, since time.Time) ([]*types.Entity, error) { args := m.Called(atype, since) return args.Get(0).([]*types.Entity), args.Error(1) @@ -474,6 +439,11 @@ func (m *mockAssetDB) CreateEntityTag(entity *types.Entity, property oam.Propert return args.Get(0).(*types.EntityTag), args.Error(1) } +func (m *mockAssetDB) FindEntityTagById(id string) (*types.EntityTag, error) { + args := m.Called(id) + return args.Get(0).(*types.EntityTag), args.Error(1) +} + func (m *mockAssetDB) GetEntityTags(entity *types.Entity, since time.Time, names ...string) ([]*types.EntityTag, error) { args := m.Called(entity, since, names) return args.Get(0).([]*types.EntityTag), args.Error(1) @@ -489,6 +459,11 @@ func (m *mockAssetDB) CreateEdgeTag(edge *types.Edge, property oam.Property) (*t return args.Get(0).(*types.EdgeTag), args.Error(1) } +func (m *mockAssetDB) FindEdgeTagById(id string) (*types.EdgeTag, error) { + args := m.Called(id) + return args.Get(0).(*types.EdgeTag), args.Error(1) +} + func (m *mockAssetDB) GetEdgeTags(edge *types.Edge, since time.Time, names ...string) ([]*types.EdgeTag, error) { args := m.Called(edge, since, names) return args.Get(0).([]*types.EdgeTag), args.Error(1) diff --git a/repository/repository.go b/repository/repository.go index cdd104e..6d11144 100644 --- a/repository/repository.go +++ b/repository/repository.go @@ -25,14 +25,15 @@ type Repository interface { FindEntityById(id string) (*types.Entity, error) FindEntityByContent(asset oam.Asset, since time.Time) ([]*types.Entity, error) FindEntitiesByType(atype oam.AssetType, since time.Time) ([]*types.Entity, error) - FindEntitiesByScope(constraints []oam.Asset, since time.Time) ([]*types.Entity, error) Link(edge *types.Edge) (*types.Edge, error) IncomingEdges(entity *types.Entity, since time.Time, labels ...string) ([]*types.Edge, error) OutgoingEdges(entity *types.Entity, since time.Time, labels ...string) ([]*types.Edge, error) CreateEntityTag(entity *types.Entity, property oam.Property) (*types.EntityTag, error) + FindEntityTagById(id string) (*types.EntityTag, error) GetEntityTags(entity *types.Entity, since time.Time, names ...string) ([]*types.EntityTag, error) DeleteEntityTag(id string) error CreateEdgeTag(edge *types.Edge, property oam.Property) (*types.EdgeTag, error) + FindEdgeTagById(id string) (*types.EdgeTag, error) GetEdgeTags(edge *types.Edge, since time.Time, names ...string) ([]*types.EdgeTag, error) DeleteEdgeTag(id string) error Close() error @@ -40,7 +41,12 @@ type Repository interface { // New creates a new instance of the asset database repository. func New(dbtype, dsn string) (Repository, error) { - if strings.EqualFold(dbtype, sqlrepo.Postgres) || strings.EqualFold(dbtype, sqlrepo.SQLite) { + switch strings.ToLower(dbtype) { + case strings.ToLower(sqlrepo.Postgres): + fallthrough + case strings.ToLower(sqlrepo.SQLite): + fallthrough + case strings.ToLower(sqlrepo.SQLiteMemory): return sqlrepo.New(dbtype, dsn) } return nil, errors.New("unknown DB type") diff --git a/repository/sqlrepo/sql_scope.go b/repository/sqlrepo/sql_scope.go deleted file mode 100644 index 06d64d2..0000000 --- a/repository/sqlrepo/sql_scope.go +++ /dev/null @@ -1,104 +0,0 @@ -// 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 sqlrepo - -import ( - "errors" - "strconv" - "time" - - "github.com/caffix/stringset" - "github.com/owasp-amass/asset-db/types" - oam "github.com/owasp-amass/open-asset-model" - "github.com/owasp-amass/open-asset-model/domain" - "gorm.io/gorm" -) - -// FindEntitiesByScope finds entities in the database by applying all the scope constraints provided and last seen after the since parameter. -// It takes a slice representing the set of constraints to serve as the scope and retrieves the corresponding entities from the database. -// If since.IsZero(), the parameter will be ignored. -// Returns a slice of matching entities as []*types.Entity or an error if the search fails. -func (sql *sqlRepository) FindEntitiesByScope(constraints []oam.Asset, since time.Time) ([]*types.Entity, error) { - var findings []*types.Entity - - for _, constraint := range constraints { - if entities, err := sql.constraintEdgeCases(constraint, since); err == nil { - for _, e := range entities { - if f, err := e.Parse(); err == nil { - findings = append(findings, &types.Entity{ - ID: strconv.FormatUint(e.ID, 10), - CreatedAt: e.CreatedAt, - LastSeen: e.LastSeen, - Asset: f, - }) - } - } - } - - if entities, err := sql.inAndOut(constraint, since); err == nil { - findings = append(findings, entities...) - } - } - - if len(findings) == 0 { - return []*types.Entity{}, errors.New("no entities in scope") - } - return findings, nil -} - -func (sql *sqlRepository) inAndOut(constraint oam.Asset, since time.Time) ([]*types.Entity, error) { - constraints, err := sql.FindEntityByContent(constraint, time.Time{}) - if err != nil || len(constraints) == 0 { - return constraints, err - } - - ids := stringset.New() - for _, constraint := range constraints { - if rels, err := sql.IncomingEdges(constraint, since); err == nil { - for _, rel := range rels { - ids.Insert(rel.FromEntity.ID) - } - } - - if rels, err := sql.OutgoingEdges(constraint, since); err == nil { - for _, rel := range rels { - ids.Insert(rel.ToEntity.ID) - } - } - } - - var entities []*types.Entity - for _, id := range ids.Slice() { - if e, err := sql.FindEntityById(id); err == nil { - entities = append(entities, e) - } - } - - if len(entities) == 0 { - return []*types.Entity{}, errors.New("no entities in scope") - } - return entities, nil -} - -func (sql *sqlRepository) constraintEdgeCases(constraint oam.Asset, since time.Time) ([]Entity, error) { - switch v := constraint.(type) { - case *domain.FQDN: - return sql.fqdnToEmails(v, since) - } - return nil, errors.New("no results found") -} - -func (sql *sqlRepository) fqdnToEmails(fqdn *domain.FQDN, since time.Time) ([]Entity, error) { - var entities []Entity - var result *gorm.DB - - if since.IsZero() { - result = sql.db.Where("etype = ? AND content->>'address' LIKE ?", oam.EmailAddress, "%"+fqdn.Name).Find(&entities) - } else { - result = sql.db.Where("etype = ? AND content->>'address' LIKE ? AND last_seen > ?", oam.EmailAddress, "%"+fqdn.Name, since).Find(&entities) - } - - return entities, result.Error -}