Skip to content

Commit

Permalink
added entity and edge tags
Browse files Browse the repository at this point in the history
  • Loading branch information
@caffix
caffix committed Nov 9, 2024
1 parent 117e6aa commit 4e7a23d
Show file tree
Hide file tree
Showing 8 changed files with 599 additions and 80 deletions.
30 changes: 30 additions & 0 deletions assetdb_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -468,3 +468,33 @@ func (m *mockAssetDB) OutgoingEdges(asset *types.Entity, since time.Time, labels
args := m.Called(asset, since, labels)
return args.Get(0).([]*types.Edge), args.Error(1)
}

func (m *mockAssetDB) CreateEntityTag(entity *types.Entity, property oam.Property) (*types.EntityTag, error) {
args := m.Called(entity, property)
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)
}

func (m *mockAssetDB) DeleteEntityTag(id string) error {
args := m.Called(id)
return args.Error(0)
}

func (m *mockAssetDB) CreateEdgeTag(edge *types.Edge, property oam.Property) (*types.EdgeTag, error) {
args := m.Called(edge, property)
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)
}

func (m *mockAssetDB) DeleteEdgeTag(id string) error {
args := m.Called(id)
return args.Error(0)
}
73 changes: 17 additions & 56 deletions repository/entity_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,7 @@ func TestMain(m *testing.M) {

func TestLastSeenUpdates(t *testing.T) {
ip, _ := netip.ParseAddr("45.73.25.1")
asset := network.IPAddress{Address: ip, Type: "IPv4"}
asset := &network.IPAddress{Address: ip, Type: "IPv4"}
a1, err := store.CreateEntity(asset)
assert.NoError(t, err)

Expand Down Expand Up @@ -240,22 +240,12 @@ func TestRepository(t *testing.T) {
for _, tc := range testCases {
t.Run(tc.description, func(t *testing.T) {
sourceEntity, err := store.CreateEntity(tc.sourceAsset)
if err != nil {
t.Fatalf("failed to create entity: %s", err)
}

if sourceEntity == nil {
t.Fatalf("failed to create entity: entity is nil")
}
assert.NoError(t, err)
assert.NotEqual(t, sourceEntity, nil)

foundAsset, err := store.FindEntityById(sourceEntity.ID, start)
if err != nil {
t.Fatalf("failed to find entity by id: %s", err)
}

if foundAsset == nil {
t.Fatalf("failed to find entity by id: found entity is nil")
}
assert.NoError(t, err)
assert.NotEqual(t, foundAsset, nil)

if foundAsset.ID != sourceEntity.ID {
t.Fatalf("failed to find entity by id: expected entity id %s, got %s", sourceEntity.ID, foundAsset.ID)
Expand All @@ -266,13 +256,8 @@ func TestRepository(t *testing.T) {
}

foundAssetByContent, err := store.FindEntityByContent(sourceEntity.Asset, start)
if err != nil {
t.Fatalf("failed to find entity by content: %s", err)
}

if foundAssetByContent == nil {
t.Fatalf("failed to find entity by content: found entity is nil")
}
assert.NoError(t, err)
assert.NotEqual(t, foundAssetByContent, nil)

if foundAssetByContent[0].ID != sourceEntity.ID {
t.Fatalf("failed to find entity by content: expected entity id %s, got %s", sourceEntity.ID, foundAssetByContent[0].ID)
Expand Down Expand Up @@ -314,13 +299,8 @@ func TestRepository(t *testing.T) {
}

destinationEntity, err := store.CreateEntity(tc.destinationAsset)
if err != nil {
t.Fatalf("failed to create destination entity: %s", err)
}

if destinationEntity == nil {
t.Fatalf("failed to create destination entity: destination entity is nil")
}
assert.NoError(t, err)
assert.NotEqual(t, destinationEntity, nil)

edge := &types.Edge{
Relation: tc.relation,
Expand All @@ -329,22 +309,12 @@ func TestRepository(t *testing.T) {
}

e, err := store.Link(edge)
if err != nil {
t.Fatalf("failed to link entities: %s", err)
}

if e == nil {
t.Fatalf("failed to link entities: edge is nil")
}
assert.NoError(t, err)
assert.NotEqual(t, e, nil)

incoming, err := store.IncomingEdges(destinationEntity, start, tc.relation.Label())
if err != nil {
t.Fatalf("failed to query incoming edges: %s", err)
}

if incoming == nil {
t.Fatalf("failed to query incoming edges: incoming edge is nil %s", err)
}
assert.NoError(t, err)
assert.NotEqual(t, incoming, nil)

if incoming[0].Relation.Label() != tc.relation.Label() {
t.Fatalf("failed to query incoming edges: expected relation %s, got %s", tc.relation, incoming[0].Relation.Label())
Expand All @@ -359,13 +329,8 @@ func TestRepository(t *testing.T) {
}

outgoing, err := store.OutgoingEdges(sourceEntity, start, tc.relation.Label())
if err != nil {
t.Fatalf("failed to query outgoing edges: %s", err)
}

if outgoing == nil {
t.Fatalf("failed to query outgoing edges: outgoing edge is nil")
}
assert.NoError(t, err)
assert.NotEqual(t, outgoing, nil)

if outgoing[0].Relation.Label() != tc.relation.Label() {
t.Fatalf("failed to query outgoing edges: expected edge %s, got %s", tc.relation, outgoing[0].Relation.Label())
Expand All @@ -380,14 +345,10 @@ func TestRepository(t *testing.T) {
}

err = store.DeleteEdge(e.ID)
if err != nil {
t.Fatalf("failed to delete edges: %s", err)
}
assert.NoError(t, err)

err = store.DeleteEntity(destinationEntity.ID)
if err != nil {
t.Fatalf("failed to delete asset: %s", err)
}
assert.NoError(t, err)

if _, err = store.FindEntityById(destinationEntity.ID, start); err == nil {
t.Fatal("failed to delete entity: the entity was not removed from the database")
Expand Down
65 changes: 55 additions & 10 deletions repository/models.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,10 @@ import (
"github.com/owasp-amass/open-asset-model/contact"
"github.com/owasp-amass/open-asset-model/domain"
oamfile "github.com/owasp-amass/open-asset-model/file"
"github.com/owasp-amass/open-asset-model/fingerprint"
"github.com/owasp-amass/open-asset-model/network"
"github.com/owasp-amass/open-asset-model/org"
"github.com/owasp-amass/open-asset-model/people"
"github.com/owasp-amass/open-asset-model/property"
oamreg "github.com/owasp-amass/open-asset-model/registration"
"github.com/owasp-amass/open-asset-model/relation"
"github.com/owasp-amass/open-asset-model/service"
Expand All @@ -34,11 +34,20 @@ type Entity struct {
Content datatypes.JSON
}

// EntityTag represents additional metadata added to an entity in the asset database.
type EntityTag struct {
ID uint64 `gorm:"primaryKey;column:tag_id"`
CreatedAt time.Time `gorm:"type:datetime;default:CURRENT_TIMESTAMP();column:created_at"`
LastSeen time.Time `gorm:"type:datetime;default:CURRENT_TIMESTAMP();column:last_seen"`
Type string `gorm:"column:ttype"`
Content datatypes.JSON
}

// Edge represents a relationship between two entities stored in the database.
type Edge struct {
ID uint64 `gorm:"primaryKey;column:edge_id"`
CreatedAt time.Time `gorm:"type:datetime;default:CURRENT_TIMESTAMP();"`
LastSeen time.Time `gorm:"type:datetime;default:CURRENT_TIMESTAMP();"`
CreatedAt time.Time `gorm:"type:datetime;default:CURRENT_TIMESTAMP();column:created_at"`
LastSeen time.Time `gorm:"type:datetime;default:CURRENT_TIMESTAMP();column:last_seen"`
Type string `gorm:"column:etype"`
Content datatypes.JSON
FromEntityID uint64 `gorm:"column:from_entity_id"`
Expand All @@ -47,6 +56,15 @@ type Edge struct {
ToEntity Entity
}

// EdgeTag represents additional metadata added to an edge in the asset database.
type EdgeTag struct {
ID uint64 `gorm:"primaryKey;column:tag_id"`
CreatedAt time.Time `gorm:"type:datetime;default:CURRENT_TIMESTAMP();column:created_at"`
LastSeen time.Time `gorm:"type:datetime;default:CURRENT_TIMESTAMP();column:last_seen"`
Type string `gorm:"column:ttype"`
Content datatypes.JSON
}

// Parse parses the content of the entity into the corresponding Open Asset Model (OAM) asset type.
// It returns the parsed asset and an error, if any.
func (e *Entity) Parse() (oam.Asset, error) {
Expand Down Expand Up @@ -89,11 +107,6 @@ func (e *Entity) Parse() (oam.Asset, error) {

err = json.Unmarshal(e.Content, &dr)
asset = &dr
case string(oam.Fingerprint):
var fingerprint fingerprint.Fingerprint

err = json.Unmarshal(e.Content, &fingerprint)
asset = &fingerprint
case string(oam.Organization):
var organization org.Organization

Expand Down Expand Up @@ -175,8 +188,6 @@ func (e *Entity) JSONQuery() (*datatypes.JSONQueryExpression, error) {
return jsonQuery.Equals(v.Handle, "handle"), nil
case *oamreg.DomainRecord:
return jsonQuery.Equals(v.Domain, "domain"), nil
case *fingerprint.Fingerprint:
return jsonQuery.Equals(v.Value, "value"), nil
case *org.Organization:
return jsonQuery.Equals(v.Name, "name"), nil
case *people.Person:
Expand Down Expand Up @@ -240,3 +251,37 @@ func (e *Edge) Parse() (oam.Relation, error) {

return rel, err
}

// Parse parses the content of the entity tag into the corresponding Open Asset Model (OAM) property type.
// It returns the parsed property and an error, if any.
func (e *EntityTag) Parse() (oam.Property, error) {
return parseProperty(e.Type, e.Content)
}

// Parse parses the content of the edge tag into the corresponding Open Asset Model (OAM) property type.
// It returns the parsed property and an error, if any.
func (e *EdgeTag) Parse() (oam.Property, error) {
return parseProperty(e.Type, e.Content)
}

func parseProperty(ptype string, content datatypes.JSON) (oam.Property, error) {
var err error
var prop oam.Property

switch ptype {
case string(oam.SimpleProperty):
var sp property.SimpleProperty

err = json.Unmarshal(content, &sp)
prop = &sp
case string(oam.VulnProperty):
var vp property.VulnProperty

err = json.Unmarshal(content, &vp)
prop = &vp
default:
return nil, fmt.Errorf("unknown property type: %s", ptype)
}

return prop, err
}
10 changes: 0 additions & 10 deletions repository/models_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ import (
oamcert "github.com/owasp-amass/open-asset-model/certificate"
"github.com/owasp-amass/open-asset-model/contact"
"github.com/owasp-amass/open-asset-model/domain"
"github.com/owasp-amass/open-asset-model/fingerprint"
"github.com/owasp-amass/open-asset-model/network"
"github.com/owasp-amass/open-asset-model/org"
"github.com/owasp-amass/open-asset-model/people"
Expand Down Expand Up @@ -87,10 +86,6 @@ func TestModels(t *testing.T) {
description: "parse location",
asset: &contact.Location{Address: "1600 Pennsylvania Ave NW, Washington, DC 20500"},
},
{
description: "parse fingerprint",
asset: &fingerprint.Fingerprint{Value: "a1:2b:3c:4d:5e:6f:7g:8h:9i:0j:1k:2l:3m:4n:5o:6p"},
},
{
description: "parse organization",
asset: &org.Organization{Name: "Example, Inc."},
Expand Down Expand Up @@ -183,11 +178,6 @@ func TestModels(t *testing.T) {
asset: &contact.Location{Address: "1600 Pennsylvania Ave NW, Washington, DC 20500"},
expectedQuery: datatypes.JSONQuery("content").Equals("1600 Pennsylvania Ave NW, Washington, DC 20500", "address"),
},
{
description: "json query for fingerprint",
asset: &fingerprint.Fingerprint{Value: "a1:2b:3c:4d:5e:6f:7g:8h:9i:0j:1k:2l:3m:4n:5o:6p"},
expectedQuery: datatypes.JSONQuery("content").Equals("a1:2b:3c:4d:5e:6f:7g:8h:9i:0j:1k:2l:3m:4n:5o:6p", "value"),
},
{
description: "json query for url",
asset: &url.URL{Raw: "https://www.example.com"},
Expand Down
10 changes: 8 additions & 2 deletions repository/repository.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,13 @@ type Repository interface {
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(asset *types.Entity, since time.Time, labels ...string) ([]*types.Edge, error)
OutgoingEdges(asset *types.Entity, since time.Time, labels ...string) ([]*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)
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)
GetEdgeTags(edge *types.Edge, since time.Time, names ...string) ([]*types.EdgeTag, error)
DeleteEdgeTag(id string) error
Close() error
}
Loading

0 comments on commit 4e7a23d

Please sign in to comment.