feat(COD-1156): support sast tables in cds

Add the parameter to specify the url type for sast tables
lacework · Sep 29, 2023 · 30456b8 · 30456b8
1 parent 95b2f48
commit 30456b8
Show file tree

Hide file tree

Showing 2 changed files with 93 additions and 3 deletions.
diff --git a/api/component_data.go b/api/component_data.go
@@ -15,11 +15,17 @@ type ComponentDataService struct {
 	client *Client
 }
 
+const URL_TYPE_DEFAULT = "Default"
+const URL_TYPE_SAST_TABLES = "SastTables"
+
+var URL_TYPES = []string{URL_TYPE_DEFAULT, URL_TYPE_SAST_TABLES}
+
 type ComponentDataInitialRequest struct {
 	Name             string          `json:"name"`
 	Tags             []string        `json:"tags"`
 	SupportedMethods []string        `json:"supportedMethods"`
 	Documents        []*DocumentSpec `json:"documents"`
+	UrlType          string          `json:"urlType"`
 }
 
 type DocumentSpec struct {
@@ -43,6 +49,7 @@ type ComponentDataUploadMethod struct {
 
 type ComponentDataCompleteRequest struct {
 	UploadGuid string `json:"uploadGuid"`
+	UrlType    string `json:"urlType"`
 }
 
 type ComponentDataCompleteResponseRaw struct {
@@ -53,8 +60,29 @@ type ComponentDataCompleteResponse struct {
 	Guid string `json:"guid,omitempty"`
 }
 
-func (svc *ComponentDataService) UploadFiles(name string, tags []string, paths []string) (string, error) {
-	initialRequest, err := buildComponentDataInitialRequest(name, tags, paths)
+func (svc *ComponentDataService) UploadFiles(
+	name string, tags []string, paths []string) (string, error) {
+	return svc.doUploadFiles(name, tags, paths, URL_TYPE_DEFAULT)
+}
+
+func (svc *ComponentDataService) UploadSastTables(
+	name string, paths []string) (string, error) {
+	return svc.doUploadFiles(name, []string{"sast"}, paths, URL_TYPE_SAST_TABLES)
+}
+
+func (svc *ComponentDataService) doUploadFiles(
+	name string, tags []string, paths []string, urlType string) (string, error) {
+	var hasValidType = false
+	for _, validType := range URL_TYPES {
+		if urlType == validType {
+			hasValidType = true
+			break
+		}
+	}
+	if !hasValidType {
+		return "", errors.Errorf("Invalid URL type: (%s)", urlType)
+	}
+	initialRequest, err := buildComponentDataInitialRequest(name, tags, paths, urlType)
 	if err != nil {
 		return "", err
 	}
@@ -88,6 +116,7 @@ func (svc *ComponentDataService) UploadFiles(name string, tags []string, paths [
 	}
 	completeRequest := ComponentDataCompleteRequest{
 		UploadGuid: initialResponse.Data.Guid,
+		UrlType:    urlType,
 	}
 	var completeResponse ComponentDataCompleteResponseRaw
 	err = doWithExponentialBackoffWaiting(func() error {
@@ -107,7 +136,7 @@ func (svc *ComponentDataService) UploadFiles(name string, tags []string, paths [
 }
 
 func buildComponentDataInitialRequest(
-	name string, tags []string, paths []string,
+	name string, tags []string, paths []string, urlType string,
 ) (*ComponentDataInitialRequest, error) {
 	documents := make([]*DocumentSpec, 0, len(paths))
 	for _, path := range paths {
@@ -125,6 +154,7 @@ func buildComponentDataInitialRequest(
 		Tags:             tags,
 		SupportedMethods: []string{"AwsS3"},
 		Documents:        documents,
+		UrlType:          urlType,
 	}, nil
 }
 

diff --git a/api/component_data_test.go b/api/component_data_test.go
@@ -42,6 +42,66 @@ func TestUploadFiles(t *testing.T) {
 	assert.Equal(t, "SOME-GUID", guid)
 }
 
+func TestDefaultUrlType(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_DEFAULT)
+		_, err := fmt.Fprintf(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_DEFAULT)
+		_, err := fmt.Fprintf(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.UploadFiles("doc-set", []string{"sast"}, []string{})
+	assert.Nil(t, err)
+	assert.Equal(t, "SOME-GUID", guid)
+}
+
+func TestSastTablesUrlType(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_SAST_TABLES)
+		_, err := fmt.Fprintf(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_SAST_TABLES)
+		_, err := fmt.Fprintf(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.UploadSastTables("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 {