refactor: use single windowSearch func

Signed-off-by: Darren Murray <[email protected]> refactor: use single windowSearch func
lacework · Oct 21, 2022 · 79acf4a · 79acf4a
1 parent be0c9c7
commit 79acf4a
Show file tree

Hide file tree

Showing 6 changed files with 75 additions and 53 deletions.
diff --git a/api/compliance_evaluations.go b/api/compliance_evaluations.go
@@ -18,6 +18,8 @@
 
 package api
 
+import "time"
+
 type ComplianceEvaluationService struct {
 	client *Client
 }
@@ -45,38 +47,23 @@ const AwsComplianceEvaluationDataset complianceEvaluationDataset = "AwsComplianc
 //  )
 //   lacework.V2.ComplianceEvaluation.Search(&awsComplianceEvaluationSearchResponse, filters)
 //
-func (svc *ComplianceEvaluationService) Search(response interface{}, filters ComplianceEvaluationSearch) error {
+func (svc *ComplianceEvaluationService) Search(response interface{}, filters SearchableFilter) error {
 	return svc.client.RequestEncoderDecoder("POST", apiV2ComplianceEvaluationsSearch, filters, response)
 }
 
-type ComplianceEvaluationSearch struct {
-	SearchFilter
-	Dataset complianceEvaluationDataset `json:"dataset"`
+func (c *ComplianceEvaluationSearch) GetTimeFilter() *TimeFilter {
+	return c.TimeFilter
 }
 
-// WindowedSearch performs a new search of a specific time frame size,
-// until response data is found or the max searchable days is reached
-func (svc *ComplianceEvaluationService) WindowedSearch(response *ComplianceEvaluationAwsResponse, filter ComplianceEvaluationSearch, size int, max int) error {
-	for i := 0; i < max; i += size {
-		err := svc.Search(&response, filter)
-		if err != nil {
-			return err
-		}
-		if len(response.Data) != 0 {
-			return nil
-		}
+func (c *ComplianceEvaluationSearch) SetStartTime(t *time.Time) {
+	c.TimeFilter.StartTime = t
+}
 
-		//adjust window
-		newStart := filter.StartTime.AddDate(0, 0, -size)
-		newEnd := filter.EndTime.AddDate(0, 0, -size)
+func (c *ComplianceEvaluationSearch) SetEndTime(t *time.Time) {
+	c.TimeFilter.EndTime = t
+}
 
-		// ensure we do not go over the max allowed searchable days
-		rem := (i - max) % size
-		if rem > 0 {
-			newEnd = filter.EndTime.AddDate(0, 0, -rem)
-		}
-		filter.TimeFilter.StartTime = &newStart
-		filter.TimeFilter.EndTime = &newEnd
-	}
-	return nil
+type ComplianceEvaluationSearch struct {
+	SearchFilter
+	Dataset complianceEvaluationDataset `json:"dataset"`
 }
diff --git a/api/compliance_evaluations_aws.go b/api/compliance_evaluations_aws.go
@@ -25,6 +25,10 @@ type ComplianceEvaluationAwsResponse struct {
 	Paging V2Pagination              `json:"paging"`
 }
 
+func (r ComplianceEvaluationAwsResponse) GetDataLength() int {
+	return len(r.Data)
+}
+
 func (r ComplianceEvaluationAwsResponse) PageInfo() *V2Pagination {
 	return &r.Paging
 }

diff --git a/api/inventory.go b/api/inventory.go
@@ -18,6 +18,8 @@
 
 package api
 
+import "time"
+
 type InventoryService struct {
 	client *Client
 }
@@ -47,7 +49,7 @@ const AwsInventoryDataset inventoryDataset = "AwsCompliance"
 //  )
 //   lacework.V2.Inventory.Search(&awsInventorySearchResponse, filters)
 //
-func (svc *InventoryService) Search(response interface{}, filters InventorySearch) error {
+func (svc *InventoryService) Search(response interface{}, filters SearchableFilter) error {
 	return svc.client.RequestEncoderDecoder("POST", apiV2InventorySearch, filters, response)
 }
 
@@ -57,29 +59,14 @@ type InventorySearch struct {
 	Dataset inventoryDataset `json:"dataset"`
 }
 
-// WindowedSearch performs a new search of a specific time frame size,
-// until response data is found or the max searchable days is reached
-func (svc *InventoryService) WindowedSearch(response *InventoryAwsResponse, filter InventorySearch, size int, max int) error {
-	for i := 0; i < max; i += size {
-		err := svc.Search(&response, filter)
-		if err != nil {
-			return err
-		}
-		if len(response.Data) != 0 {
-			return nil
-		}
+func (i InventorySearch) GetTimeFilter() *TimeFilter {
+	return i.TimeFilter
+}
 
-		//adjust window
-		newStart := filter.StartTime.AddDate(0, 0, -size)
-		newEnd := filter.EndTime.AddDate(0, 0, -size)
+func (i InventorySearch) SetStartTime(time *time.Time) {
+	i.TimeFilter.StartTime = time
+}
 
-		// ensure we do not go over the max allowed searchable days
-		rem := (i - max) % size
-		if rem > 0 {
-			newEnd = filter.EndTime.AddDate(0, 0, -rem)
-		}
-		filter.TimeFilter.StartTime = &newStart
-		filter.TimeFilter.EndTime = &newEnd
-	}
-	return nil
+func (i InventorySearch) SetEndTime(time *time.Time) {
+	i.TimeFilter.EndTime = time
 }
diff --git a/api/inventory_aws.go b/api/inventory_aws.go
@@ -23,6 +23,10 @@ type InventoryAwsResponse struct {
 	Paging V2Pagination   `json:"paging"`
 }
 
+func (r InventoryAwsResponse) GetDataLength() int {
+	return len(r.Data)
+}
+
 func (r InventoryAwsResponse) PageInfo() *V2Pagination {
 	return &r.Paging
 }

diff --git a/api/v2_search_filters.go b/api/v2_search_filters.go
@@ -54,3 +54,42 @@ type TimeFilter struct {
 	StartTime *time.Time `json:"startTime,omitempty"`
 	EndTime   *time.Time `json:"endTime,omitempty"`
 }
+
+type SearchResponse interface {
+	GetDataLength() int
+}
+
+type SearchableFilter interface {
+	GetTimeFilter() *TimeFilter
+	SetStartTime(*time.Time)
+	SetEndTime(*time.Time)
+}
+
+type search func(response interface{}, filters SearchableFilter) error
+
+// WindowedSearch performs a new search of a specific time frame size,
+// until response data is found or the max searchable days is reached
+func WindowedSearch(fn search, size int, max int, response SearchResponse, filter SearchableFilter) error {
+	for i := 0; i < max; i += size {
+		err := fn(&response, filter)
+		if err != nil {
+			return err
+		}
+		if response.GetDataLength() != 0 {
+			return nil
+		}
+
+		//adjust window
+		newStart := filter.GetTimeFilter().StartTime.AddDate(0, 0, -size)
+		newEnd := filter.GetTimeFilter().EndTime.AddDate(0, 0, -size)
+
+		// ensure we do not go over the max allowed searchable days
+		rem := (i - max) % size
+		if rem > 0 {
+			newEnd = filter.GetTimeFilter().EndTime.AddDate(0, 0, -rem)
+		}
+		filter.SetStartTime(&newStart)
+		filter.SetEndTime(&newEnd)
+	}
+	return nil
+}
diff --git a/cli/cmd/compliance_aws.go b/cli/cmd/compliance_aws.go
@@ -448,7 +448,7 @@ The output from status with the --json flag can be used in the body of PATCH api
 					Csp:     api.AwsInventoryType,
 				}
 			)
-			err := cli.LwApi.V2.Inventory.WindowedSearch(&awsInventorySearchResponse, filter, 7, 92)
+			err := api.WindowedSearch(cli.LwApi.V2.Inventory.Search, 7, 92, &awsInventorySearchResponse, &filter)
 			cli.StopProgress()
 
 			if len(awsInventorySearchResponse.Data) == 0 {
@@ -479,7 +479,8 @@ The output from status with the --json flag can be used in the body of PATCH api
 					Dataset: api.AwsComplianceEvaluationDataset,
 				}
 			)
-			err = cli.LwApi.V2.ComplianceEvaluations.WindowedSearch(&awsComplianceEvaluationSearchResponse, complianceFilter, 7, 92)
+
+			err = api.WindowedSearch(cli.LwApi.V2.ComplianceEvaluations.Search, 7, 92, &awsComplianceEvaluationSearchResponse, &complianceFilter)
 			cli.StopProgress()
 			if err != nil {
 				return err