feat: add incoming and outgoing relations to repository (#8)

* feat: add incoming and outgoing relations to respository --------- Co-authored-by: Felipe Ruiz <fruiz.rob@gmail.com>
owasp-amass · Jun 7, 2023 · 260d3b0 · 260d3b0
1 parent deb725f
commit 260d3b0
Show file tree

Hide file tree

Showing 6 changed files with 337 additions and 20 deletions.
diff --git a/Makefile b/Makefile
@@ -5,4 +5,4 @@ include .env.local
 export $(shell sed 's/=.*//' .env.local)
 
 test:
-	go test -v ./...
+	# go test -v -cover ./...
diff --git a/assetdb.go b/assetdb.go
@@ -32,12 +32,12 @@ func (as *assetDB) Create(source *types.Asset, relation *string, discovered oam.
 
 	newAsset, err := as.repository.CreateAsset(discovered)
 	if err != nil {
-		return &types.Asset{}, err
+		return nil, err
 	}
 
 	_, err = as.repository.Link(source, *relation, newAsset)
 	if err != nil {
-		return &types.Asset{}, err
+		return nil, err
 	}
 
 	return newAsset, nil
@@ -54,3 +54,15 @@ func (as *assetDB) FindByContent(asset oam.Asset) ([]*types.Asset, error) {
 func (as *assetDB) FindById(id string) (*types.Asset, error) {
 	return as.repository.FindAssetById(id)
 }
+
+// IncomingRelations finds all relations pointing to `asset“ for the specified `relationTypes`, if any.
+// If no `relationTypes` are specified, all incoming relations are returned.
+func (as *assetDB) IncomingRelations(asset *types.Asset, relationTypes ...string) ([]*types.Relation, error) {
+	return as.repository.IncomingRelations(asset, relationTypes...)
+}
+
+// OutgoingRelations finds all relations from `asset“ to another asset for the specified `relationTypes`, if any.
+// If no `relationTypes` are specified, all outgoing relations are returned.
+func (as *assetDB) OutgoingRelations(asset *types.Asset, relationTypes ...string) ([]*types.Relation, error) {
+	return as.repository.OutgoingRelations(asset, relationTypes...)
+}
diff --git a/assetdb_test.go b/assetdb_test.go
@@ -38,6 +38,16 @@ func (m *mockAssetDB) Link(source *types.Asset, relation string, destination *ty
 	return args.Get(0).(*types.Relation), args.Error(1)
 }
 
+func (m *mockAssetDB) IncomingRelations(asset *types.Asset, relationTypes ...string) ([]*types.Relation, error) {
+	args := m.Called(asset, relationTypes)
+	return args.Get(0).([]*types.Relation), args.Error(1)
+}
+
+func (m *mockAssetDB) OutgoingRelations(asset *types.Asset, relationTypes ...string) ([]*types.Relation, error) {
+	args := m.Called(asset, relationTypes)
+	return args.Get(0).([]*types.Relation), args.Error(1)
+}
+
 func TestMain(m *testing.M) {
 	exitVal := m.Run()
 
@@ -106,8 +116,8 @@ func TestAssetDB(t *testing.T) {
 			expected      *types.Asset
 			expectedError error
 		}{
-			{"An asset is found", "1", &types.Asset{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}}, nil},
-			{"An asset is not found", "2", &types.Asset{}, fmt.Errorf("Asset not found")},
+			{"an asset is found", "1", &types.Asset{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}}, nil},
+			{"an asset is not found", "2", &types.Asset{}, fmt.Errorf("Asset not found")},
 		}
 
 		for _, tc := range testCases {
@@ -136,8 +146,8 @@ func TestAssetDB(t *testing.T) {
 			expected      []*types.Asset
 			expectedError error
 		}{
-			{"An asset is found", domain.FQDN{Name: "www.domain.com"}, []*types.Asset{{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}}}, nil},
-			{"An asset is not found", domain.FQDN{Name: "www.domain.com"}, []*types.Asset{}, fmt.Errorf("Asset not found")},
+			{"an asset is found", domain.FQDN{Name: "www.domain.com"}, []*types.Asset{{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}}}, nil},
+			{"an asset is not found", domain.FQDN{Name: "www.domain.com"}, []*types.Asset{}, fmt.Errorf("Asset not found")},
 		}
 
 		for _, tc := range testCases {
@@ -158,4 +168,116 @@ func TestAssetDB(t *testing.T) {
 			})
 		}
 	})
+
+	t.Run("IncomingRelations", func(t *testing.T) {
+		testCases := []struct {
+			description   string
+			asset         *types.Asset
+			relationTypes []string
+			expected      []*types.Relation
+			expectedError error
+		}{
+			{
+				description:   "successfully find incoming relations",
+				asset:         &types.Asset{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}},
+				relationTypes: []string{"ns_record", "cname_record"},
+				expected: []*types.Relation{
+					{
+						ID:        "1",
+						Type:      "ns_record",
+						FromAsset: &types.Asset{ID: "2", Asset: domain.FQDN{Name: "www.subdomain1.com"}},
+						ToAsset:   &types.Asset{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}},
+					},
+					{
+						ID:        "2",
+						Type:      "cname_record",
+						FromAsset: &types.Asset{ID: "3", Asset: domain.FQDN{Name: "www.subdomain2.com"}},
+						ToAsset:   &types.Asset{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}},
+					},
+				},
+				expectedError: nil,
+			},
+			{
+				description:   "error finding incoming relations",
+				asset:         &types.Asset{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}},
+				relationTypes: []string{"ns_record", "cname_record"},
+				expected:      []*types.Relation{},
+				expectedError: fmt.Errorf("error finding incoming relations"),
+			},
+		}
+
+		for _, tc := range testCases {
+			t.Run(tc.description, func(t *testing.T) {
+				mockAssetDB := new(mockAssetDB)
+				adb := assetDB{
+					repository: mockAssetDB,
+				}
+
+				mockAssetDB.On("IncomingRelations", tc.asset, tc.relationTypes).Return(tc.expected, tc.expectedError)
+
+				result, err := adb.IncomingRelations(tc.asset, tc.relationTypes...)
+
+				assert.Equal(t, tc.expected, result)
+				assert.Equal(t, tc.expectedError, err)
+
+				mockAssetDB.AssertExpectations(t)
+			})
+		}
+	})
+
+	t.Run("OutgoingRelations", func(t *testing.T) {
+		testCases := []struct {
+			description   string
+			asset         *types.Asset
+			relationTypes []string
+			expected      []*types.Relation
+			expectedError error
+		}{
+			{
+				description:   "successfully find incoming relations",
+				asset:         &types.Asset{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}},
+				relationTypes: []string{"ns_record", "cname_record"},
+				expected: []*types.Relation{
+					{
+						ID:        "1",
+						Type:      "ns_record",
+						FromAsset: &types.Asset{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}},
+						ToAsset:   &types.Asset{ID: "2", Asset: domain.FQDN{Name: "www.subdomain1.com"}},
+					},
+					{
+						ID:        "2",
+						Type:      "cname_record",
+						FromAsset: &types.Asset{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}},
+						ToAsset:   &types.Asset{ID: "2", Asset: domain.FQDN{Name: "www.subdomain2.com"}},
+					},
+				},
+				expectedError: nil,
+			},
+			{
+				description:   "error finding outgoing relations",
+				asset:         &types.Asset{ID: "1", Asset: domain.FQDN{Name: "www.domain.com"}},
+				relationTypes: []string{"ns_record", "cname_record"},
+				expected:      []*types.Relation{},
+				expectedError: fmt.Errorf("error finding outgoing relations"),
+			},
+		}
+
+		for _, tc := range testCases {
+			t.Run(tc.description, func(t *testing.T) {
+				mockAssetDB := new(mockAssetDB)
+				adb := assetDB{
+					repository: mockAssetDB,
+				}
+
+				mockAssetDB.On("OutgoingRelations", tc.asset, tc.relationTypes).Return(tc.expected, tc.expectedError)
+
+				result, err := adb.OutgoingRelations(tc.asset, tc.relationTypes...)
+
+				assert.Equal(t, tc.expected, result)
+				assert.Equal(t, tc.expectedError, err)
+
+				mockAssetDB.AssertExpectations(t)
+			})
+		}
+	})
 }
diff --git a/repository/repository.go b/repository/repository.go
@@ -14,4 +14,6 @@ type Repository interface {
 	FindAssetById(id string) (*types.Asset, error)
 	FindAssetByContent(asset oam.Asset) ([]*types.Asset, error)
 	Link(source *types.Asset, relation string, destination *types.Asset) (*types.Relation, error)
+	IncomingRelations(asset *types.Asset, relationTypes ...string) ([]*types.Relation, error)
+	OutgoingRelations(asset *types.Asset, relationTypes ...string) ([]*types.Relation, error)
 }
diff --git a/repository/sql.go b/repository/sql.go
@@ -188,3 +188,80 @@ func (sql *sqlRepository) Link(source *types.Asset, relation string, destination
 		ToAsset:   destination,
 	}, nil
 }
+
+// IncomingRelations finds all relations pointing to the asset for the specified relation types, if any.
+// If no relationTypes are specified, all outgoing relations are returned.
+func (sql *sqlRepository) IncomingRelations(asset *types.Asset, relationTypes ...string) ([]*types.Relation, error) {
+	assetId, err := strconv.ParseInt(asset.ID, 10, 64)
+	if err != nil {
+		return nil, err
+	}
+
+	relations := []Relation{}
+	if len(relationTypes) > 0 {
+		res := sql.db.Where("to_asset_id = ? AND type IN ?", assetId, relationTypes).Find(&relations)
+		if res.Error != nil {
+			return nil, res.Error
+		}
+	} else {
+		res := sql.db.Where("to_asset_id = ?", assetId).Find(&relations)
+		if res.Error != nil {
+			return nil, res.Error
+		}
+	}
+
+	return toRelations(relations), nil
+}
+
+// OutgoingRelations finds all relations from the asset to another asset for the specified relation types, if any.
+// If no relationTypes are specified, all outgoing relations are returned.
+func (sql *sqlRepository) OutgoingRelations(asset *types.Asset, relationTypes ...string) ([]*types.Relation, error) {
+	assetId, err := strconv.ParseInt(asset.ID, 10, 64)
+	if err != nil {
+		return nil, err
+	}
+
+	relations := []Relation{}
+
+	if len(relationTypes) > 0 {
+		res := sql.db.Where("from_asset_id = ? AND type IN ?", assetId, relationTypes).Find(&relations)
+		if res.Error != nil {
+			return nil, res.Error
+		}
+	} else {
+		res := sql.db.Where("from_asset_id = ?", assetId).Find(&relations)
+		if res.Error != nil {
+			return nil, res.Error
+		}
+	}
+
+	return toRelations(relations), nil
+}
+
+// toRelation converts a database Relation to a types.Relation.
+func toRelation(r Relation) *types.Relation {
+	rel := &types.Relation{
+		ID:   strconv.FormatInt(r.ID, 10),
+		Type: r.Type,
+		FromAsset: &types.Asset{
+			ID: strconv.FormatInt(r.FromAssetID, 10),
+			// Not joining to Asset to get Content
+		},
+		ToAsset: &types.Asset{
+			ID: strconv.FormatInt(r.ToAssetID, 10),
+			// Not joining to Asset to get Content
+		},
+	}
+
+	return rel
+}
+
+// toRelations converts a slice database Relations to a slice of types.Relation structs.
+func toRelations(relations []Relation) []*types.Relation {
+	var res []*types.Relation
+	for _, r := range relations {
+		res = append(res, toRelation(r))
+	}
+
+	return res
+}