Fixed testcase TestGuardDutyBenchmark

Author: akhil-sumologic

cloudwatchevents/test/test-guardduty-benchmark.py

@@ -10,6 +10,8 @@
 import boto3
 from sumologic import SumoLogic
 
+TIMEFORMAT = '%Y-%m-%dT%H:%M:%S'
+
 # Update the below values in case the template locations are changed.
 
 GUARD_DUTY_BENCHMARK_TEMPLATE = "guarddutybenchmark/template_v2.yaml"
@@ -153,10 +155,18 @@ def delete_source(self, collector_id, source):
     def fetch_logs(self):
         raw_messages = []
         # fetch Last 10 Minutes logs
-        to_time = int(time.time()) * 1000
-        from_time = to_time - self.delay * 60 * 1000
+        # Get the current time
+        to_time = datetime.datetime.now()
+        from_time = to_time - datetime.timedelta(minutes=self.delay+2)
+        from_time = from_time.strftime(TIMEFORMAT)
+        to_time = to_time.strftime(TIMEFORMAT)
+        print("Fetching records")
+        print("from_time", from_time)
+        print("to_time", to_time)
+
         search_query = '_sourceCategory=%s' % self.source_category
-        search_job_response = self.sumo.search_job(search_query, fromTime=from_time, toTime=to_time, timeZone="IST")
+        search_job_response = self.sumo.search_job(search_query, fromTime=from_time,
+                                                   toTime=to_time, timeZone="IST")
         print("Search Jobs API success with JOB ID as %s." % search_job_response["id"])
         state = "GATHERING RESULTS"
         message_count = 0
@@ -180,9 +190,25 @@ def fetch_logs(self):
         print("Received message count as %s." % len(raw_messages))
         return raw_messages
 
+    def fetch_logs_with_retry(self, retries=6, delay=60):
+        for attempt in range(1, retries + 1):
+            result = self.fetch_logs()
+            if len(result) >= 3:
+                return result
+            else:
+                print(f"Attempt {attempt}")
+                if attempt == retries:
+                    print("All attempts failed.")
+                    return result
+                else:
+                    print(f"Retrying in {delay} seconds...")
+                    self.delay += 1
+                    time.sleep(delay)
+
     # Validate the specific findings generated
     def assert_logs(self):
-        messages = self.fetch_logs()
+        messages = self.fetch_logs_with_retry()
+
         for finding_type in self.findings:
             try:
                 assert any((("type" in d and d["type"] == finding_type)
@@ -300,7 +326,7 @@ def setUp(self):
         self.source_name = "GuardDuty Benchmark"
         self.source_category = "Labs/test/guard/duty/benchmark"
         self.finding_types = ["Policy:S3/AccountBlockPublicAccessDisabled", "Policy:S3/BucketPublicAccessGranted"]
-        self.delay = 7
+        self.delay = 4
 
         # Get GuardDuty details
         self.guard_duty = boto3.client('guardduty', AWS_REGION)