From fef3bcc144ce2b661d89233650becb1a73a11a73 Mon Sep 17 00:00:00 2001
From: AlvoBen <alvo@post.bgu.ac.il>
Date: Sun, 14 Jul 2024 10:05:10 +0300
Subject: [PATCH] refactor

---
 internal/services/export.go | 71 ++++++++++++++++++++++++++-----------
 1 file changed, 50 insertions(+), 21 deletions(-)

diff --git a/internal/services/export.go b/internal/services/export.go
index 3bebf3bc4..88da9f408 100644
--- a/internal/services/export.go
+++ b/internal/services/export.go
@@ -17,45 +17,74 @@ const (
 	completedStatus     = "completed"
 )
 
-func ExportSbomResults(exportWrapper wrappers.ExportWrapper,
+func ExportSbomResults(
+	exportWrapper wrappers.ExportWrapper,
 	targetFile string,
 	results *wrappers.ResultSummary,
-	formatSbomOptions string) error {
+	formatSbomOptions string,
+) error {
+	payload, err := preparePayload(results.ScanID, formatSbomOptions)
+	if err != nil {
+		return err
+	}
+
+	log.Println("Generating SBOM report with " + payload.FileFormat + " file format")
+	sbomresp, err := exportWrapper.GenerateSbomReport(payload)
+	if err != nil {
+		return err
+	}
+
+	pollingResp, err := pollForCompletion(exportWrapper, sbomresp.ExportID)
+	if err != nil {
+		return err
+	}
+
+	if err := exportWrapper.DownloadSbomReport(pollingResp.ExportID, targetFile); err != nil {
+		return errors.Wrapf(err, "Failed downloading SBOM report")
+	}
+	return nil
+}
+
+func preparePayload(scanID, formatSbomOptions string) (*wrappers.ExportRequestPayload, error) {
 	payload := &wrappers.ExportRequestPayload{
-		ScanID:     results.ScanID,
+		ScanID:     scanID,
 		FileFormat: DefaultSbomOption,
 	}
+
 	if formatSbomOptions != "" && formatSbomOptions != DefaultSbomOption {
 		format, err := validateSbomOptions(formatSbomOptions)
 		if err != nil {
-			return err
+			return nil, err
 		}
 		payload.FileFormat = format
 	}
-	pollingResp := &wrappers.ExportPollingResponse{}
 
-	sbomresp, err := exportWrapper.GenerateSbomReport(payload)
-	if err != nil {
-		return err
-	}
+	return payload, nil
+}
+
+func pollForCompletion(exportWrapper wrappers.ExportWrapper, exportID string) (*wrappers.ExportPollingResponse, error) {
+	timeout := time.After(5 * time.Minute)
+	pollingResp := &wrappers.ExportPollingResponse{ExportStatus: exportingStatus}
 
-	log.Println("Generating SBOM report with " + payload.FileFormat + " file format")
-	pollingResp.ExportStatus = exportingStatus
 	for pollingResp.ExportStatus == exportingStatus || pollingResp.ExportStatus == pendingStatus {
-		pollingResp, err = exportWrapper.GetSbomReportStatus(sbomresp.ExportID)
-		if err != nil {
-			return errors.Wrapf(err, "%s", "failed getting SBOM report status")
+		select {
+		case <-timeout:
+			return nil, errors.Errorf("SBOM generating failed - Timed out after 5 minutes")
+		default:
+			resp, err := exportWrapper.GetSbomReportStatus(exportID)
+			if err != nil {
+				return nil, errors.Wrapf(err, "failed getting SBOM report status")
+			}
+			pollingResp = resp
+			time.Sleep(delayValueForReport * time.Second)
 		}
-		time.Sleep(delayValueForReport * time.Second)
 	}
+
 	if !strings.EqualFold(pollingResp.ExportStatus, completedStatus) {
-		return errors.Errorf("SBOM generating failed - Current status: %s", pollingResp.ExportStatus)
+		return nil, errors.Errorf("SBOM generating failed - Current status: %s", pollingResp.ExportStatus)
 	}
-	err = exportWrapper.DownloadSbomReport(pollingResp.ExportID, targetFile)
-	if err != nil {
-		return errors.Wrapf(err, "%s", "Failed downloading SBOM report")
-	}
-	return nil
+
+	return pollingResp, nil
 }
 
 func validateSbomOptions(sbomOption string) (string, error) {