From 33554aba077a583aa51b03791b52385775d6e9c4 Mon Sep 17 00:00:00 2001
From: jon-stewart <jon.stewart@lacework.net>
Date: Thu, 6 Jun 2024 14:40:14 +0100
Subject: [PATCH] feat: prosast component data upload

---
 api/component_data.go      |  7 ++++++-
 api/component_data_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/api/component_data.go b/api/component_data.go
index 6ceb15157..b434d69e5 100644
--- a/api/component_data.go
+++ b/api/component_data.go
@@ -17,8 +17,9 @@ type ComponentDataService struct {
 
 const URL_TYPE_DEFAULT = "Default"
 const URL_TYPE_SAST_TABLES = "SastTables"
+const URL_TYPE_PROSAST = "ProSast"
 
-var URL_TYPES = []string{URL_TYPE_DEFAULT, URL_TYPE_SAST_TABLES}
+var URL_TYPES = []string{URL_TYPE_DEFAULT, URL_TYPE_SAST_TABLES, URL_TYPE_PROSAST}
 
 type ComponentDataInitialRequest struct {
 	Name             string          `json:"name"`
@@ -70,6 +71,10 @@ func (svc *ComponentDataService) UploadSastTables(
 	return svc.doUploadFiles(name, []string{"sast"}, paths, URL_TYPE_SAST_TABLES)
 }
 
+func (svc *ComponentDataService) UploadProSast(name string, paths []string) (string, error) {
+	return svc.doUploadFiles(name, []string{"sast"}, paths, URL_TYPE_PROSAST)
+}
+
 func (svc *ComponentDataService) doUploadFiles(
 	name string, tags []string, paths []string, urlType string) (string, error) {
 	var hasValidType = false
diff --git a/api/component_data_test.go b/api/component_data_test.go
index 669b7f277..8074e6b8d 100644
--- a/api/component_data_test.go
+++ b/api/component_data_test.go
@@ -102,6 +102,36 @@ func TestSastTablesUrlType(t *testing.T) {
 	assert.Equal(t, "SOME-GUID", guid)
 }
 
+func TestProSastUrlType(t *testing.T) {
+	fakeServer := lacework.MockServer()
+	fakeServer.MockToken("TOKEN")
+	defer fakeServer.Close()
+	fakeServer.MockAPI("ComponentData/requestUpload", func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "POST", r.Method)
+		assert.NotNil(t, r.Body)
+		body := httpBodySniffer(r)
+		assert.Contains(t, body, api.URL_TYPE_PROSAST)
+		_, err := fmt.Fprint(w, generateInitialResponse())
+		assert.Nil(t, err)
+	})
+	fakeServer.MockAPI("ComponentData/completeUpload", func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "POST", r.Method)
+		assert.NotNil(t, r.Body)
+		body := httpBodySniffer(r)
+		assert.Contains(t, body, api.URL_TYPE_PROSAST)
+		_, err := fmt.Fprint(w, generateCompleteResponse())
+		assert.Nil(t, err)
+	})
+	c, err := api.NewClient("test",
+		api.WithToken("TOKEN"),
+		api.WithURL(fakeServer.URL()),
+	)
+	assert.Nil(t, err)
+	guid, err := c.V2.ComponentData.UploadProSast("doc-set", []string{})
+	assert.Nil(t, err)
+	assert.Equal(t, "SOME-GUID", guid)
+}
+
 func TestDoWithExponentialBackoffAlwaysFailing(t *testing.T) {
 	waited := 0
 	err := api.DoWithExponentialBackoff(func() error {