Fix vulnerability publish on broker

ZupIT · May 8, 2021 · 36736ee · 36736ee
1 parent 4b2cce3
commit 36736ee
Show file tree

Hide file tree

Showing 2 changed files with 99 additions and 5 deletions.
diff --git a/vulnerability/internal/repositories/vulnerability/vulnerability.go b/vulnerability/internal/repositories/vulnerability/vulnerability.go
@@ -3,6 +3,8 @@ package vulnerability
 import (
 	"fmt"
 
+	"github.com/ZupIT/horusec-devkit/pkg/services/database/enums"
+
 	"github.com/ZupIT/horusec-devkit/pkg/entities/analysis"
 
 	"github.com/ZupIT/horusec-devkit/pkg/enums/severities"
@@ -41,12 +43,11 @@ func NewRepositoryVulnerability(connection *database.Connection) IRepositoryVuln
 
 func (r *RepositoryVulnerability) GetAllAnalysisByVulnerabilityID(
 	vulnerabilityID uuid.UUID) ([]analysis.Analysis, error) {
-	preloads := map[string][]interface{}{
-		"AnalysisVulnerabilities":               {},
-		"AnalysisVulnerabilities.Vulnerability": {map[string]interface{}{"vulnerability_id": vulnerabilityID}},
+	condition, preloads, err := r.getConditionAndPreloadsToGetAllAnalysisByVulnerabilityID(vulnerabilityID)
+	if err != nil {
+		return nil, err
 	}
-	res := r.databaseRead.FindPreload(&[]analysis.Analysis{},
-		map[string]interface{}{}, preloads, (&analysis.Analysis{}).GetTable())
+	res := r.databaseRead.FindPreload(&[]analysis.Analysis{}, condition, preloads, (&analysis.Analysis{}).GetTable())
 	if res.GetError() != nil {
 		return nil, res.GetError()
 	}
@@ -56,6 +57,62 @@ func (r *RepositoryVulnerability) GetAllAnalysisByVulnerabilityID(
 	return *res.GetData().(*[]analysis.Analysis), nil
 }
 
+func (r *RepositoryVulnerability) getConditionAndPreloadsToGetAllAnalysisByVulnerabilityID(
+	vulnerabilityID uuid.UUID) (map[string]interface{}, map[string][]interface{}, error) {
+	analysisIDs, err := r.getAllAnalysisEntityByVulnerabilityID(vulnerabilityID)
+	if err != nil {
+		return nil, nil, err
+	}
+	preloads := map[string][]interface{}{
+		"AnalysisVulnerabilities":               {},
+		"AnalysisVulnerabilities.Vulnerability": {},
+	}
+	condition := map[string]interface{}{
+		"analysis_id": analysisIDs,
+	}
+	return condition, preloads, nil
+}
+
+func (r *RepositoryVulnerability) getAllAnalysisEntityByVulnerabilityID(
+	vulnerabilityID uuid.UUID) ([]uuid.UUID, error) {
+	entity := &[]analysis.Analysis{}
+	rawSQL := `SELECT analysis.* FROM analysis
+		JOIN analysis_vulnerabilities ON analysis.analysis_id = analysis_vulnerabilities.analysis_id
+		JOIN vulnerabilities ON vulnerabilities.vulnerability_id = analysis_vulnerabilities.vulnerability_id
+		WHERE vulnerabilities.vulnerability_id = ?`
+	res := r.databaseRead.Raw(rawSQL, entity, vulnerabilityID)
+	if res.GetError() != nil {
+		return nil, res.GetError()
+	}
+	if res.GetData() == nil {
+		return nil, enums.ErrorNotFoundRecords
+	}
+	return r.filterAnalysisByRepository(*res.GetData().(*[]analysis.Analysis)), nil
+}
+
+// nolint:funlen,gocyclo // method is not necessary broker
+func (r *RepositoryVulnerability) filterAnalysisByRepository(listExists []analysis.Analysis) (response []uuid.UUID) {
+	entities := []analysis.Analysis{}
+	for keyExists := range listExists {
+		exists := false
+		for keyEntity := range entities {
+			if listExists[keyExists].RepositoryID == entities[keyEntity].RepositoryID {
+				exists = true
+				break
+			}
+		}
+		if !exists {
+			entities = append(entities, analysis.Analysis{
+				ID: listExists[keyExists].ID, RepositoryID: listExists[keyExists].RepositoryID,
+			})
+		}
+	}
+	for key := range entities {
+		response = append(response, entities[key].ID)
+	}
+	return response
+}
+
 func (r *RepositoryVulnerability) UpdateVulnerability(
 	updateVulnerability *entityVulnerability.UpdateVulnerability) error {
 	condition := map[string]interface{}{"vulnerability_id": updateVulnerability.VulnerabilityID}

diff --git a/vulnerability/internal/repositories/vulnerability/vulnerability_test.go b/vulnerability/internal/repositories/vulnerability/vulnerability_test.go
@@ -158,6 +158,11 @@ func TestRepositoryVulnerability_UpdateVulnerability(t *testing.T) {
 func TestRepositoryVulnerability_GetVulnerabilityBYID(t *testing.T) {
 	t.Run("Should get vulnerability by id without error", func(t *testing.T) {
 		dbReadMock := &database.Mock{}
+		repositoryID := uuid.New()
+		dbReadMock.On("Raw").Return(response.NewResponse(0, nil, &[]analysis.Analysis{
+			{ID: uuid.New(), RepositoryID: repositoryID},
+			{ID: uuid.New(), RepositoryID: repositoryID},
+		}))
 		dbReadMock.On("FindPreload").Return(response.NewResponse(0, nil, &[]analysis.Analysis{
 			{ID: uuid.New()},
 		}))
@@ -168,8 +173,37 @@ func TestRepositoryVulnerability_GetVulnerabilityBYID(t *testing.T) {
 		assert.NotEmpty(t, analysisFound)
 		assert.NoError(t, err)
 	})
+	t.Run("Should get vulnerability by id with error on get analysis ID", func(t *testing.T) {
+		dbReadMock := &database.Mock{}
+		dbReadMock.On("Raw").Return(response.NewResponse(0, errors.New("unexpected error"), nil))
+		dbReadMock.On("FindPreload").Return(response.NewResponse(0, nil, &[]analysis.Analysis{
+			{ID: uuid.New()},
+		}))
+		connectionMock := &database.Connection{
+			Read: dbReadMock,
+		}
+		analysisFound, err := NewRepositoryVulnerability(connectionMock).GetAllAnalysisByVulnerabilityID(uuid.New())
+		assert.Empty(t, analysisFound)
+		assert.Error(t, err)
+	})
+	t.Run("Should get vulnerability by id with error data empty on get analysis ID", func(t *testing.T) {
+		dbReadMock := &database.Mock{}
+		dbReadMock.On("Raw").Return(response.NewResponse(0, nil, nil))
+		dbReadMock.On("FindPreload").Return(response.NewResponse(0, nil, &[]analysis.Analysis{
+			{ID: uuid.New()},
+		}))
+		connectionMock := &database.Connection{
+			Read: dbReadMock,
+		}
+		analysisFound, err := NewRepositoryVulnerability(connectionMock).GetAllAnalysisByVulnerabilityID(uuid.New())
+		assert.Empty(t, analysisFound)
+		assert.Error(t, err)
+	})
 	t.Run("Should get vulnerability by id with error", func(t *testing.T) {
 		dbReadMock := &database.Mock{}
+		dbReadMock.On("Raw").Return(response.NewResponse(0, nil, &[]analysis.Analysis{
+			{ID: uuid.New()},
+		}))
 		dbReadMock.On("FindPreload").Return(response.NewResponse(0, errors.New("unexpected error"), nil))
 		connectionMock := &database.Connection{
 			Read: dbReadMock,
@@ -180,6 +214,9 @@ func TestRepositoryVulnerability_GetVulnerabilityBYID(t *testing.T) {
 	})
 	t.Run("Should get vulnerability by id without error but data is empty", func(t *testing.T) {
 		dbReadMock := &database.Mock{}
+		dbReadMock.On("Raw").Return(response.NewResponse(0, nil, &[]analysis.Analysis{
+			{ID: uuid.New()},
+		}))
 		dbReadMock.On("FindPreload").Return(response.NewResponse(0, nil, nil))
 		connectionMock := &database.Connection{
 			Read: dbReadMock,