diff --git a/scan.py b/scan.py index 48545a06..49ba18c0 100644 --- a/scan.py +++ b/scan.py @@ -48,6 +48,9 @@ def event_object(event, event_source="s3"): # SNS events are slightly different if event_source.upper() == "SNS": event = json.loads(event["Records"][0]["Sns"]["Message"]) + # SQS events are slightly different + elif event_source.upper() == "SQS": + event = json.loads(event["Records"][0]["body"]) # Break down the record records = event["Records"] diff --git a/scan_test.py b/scan_test.py index 0a18132a..3d333057 100644 --- a/scan_test.py +++ b/scan_test.py @@ -66,6 +66,22 @@ def test_sns_event_object(self): expected_s3_object = self.s3.Object(self.s3_bucket_name, self.s3_key_name) self.assertEquals(s3_obj, expected_s3_object) + def test_sqs_event_object(self): + event = { + "Records": [ + { + "s3": { + "bucket": {"name": self.s3_bucket_name}, + "object": {"key": self.s3_key_name}, + } + } + ] + } + sqs_event = {"Records": [{"body": json.dumps(event)}]} + s3_obj = event_object(sqs_event, event_source="sqs") + expected_s3_object = self.s3.Object(self.s3_bucket_name, self.s3_key_name) + self.assertEquals(s3_obj, expected_s3_object) + def test_s3_event_object(self): event = { "Records": [