From 25d1213407eff5e2197409ca5e7069f70bc2d406 Mon Sep 17 00:00:00 2001
From: 1101-1 <70093559+1101-1@users.noreply.github.com>
Date: Wed, 2 Oct 2024 22:32:27 +0500
Subject: [PATCH] [aws][feat] Add lifecycle policy to the S3 resource (#2220)

---
 plugins/aws/fix_plugin_aws/resource/ecr.py |  3 ++-
 plugins/aws/fix_plugin_aws/resource/s3.py  | 15 +++++++++++++++
 plugins/aws/test/resources/s3_test.py      |  6 +++---
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/plugins/aws/fix_plugin_aws/resource/ecr.py b/plugins/aws/fix_plugin_aws/resource/ecr.py
index 81bc06ac20..b8ad1f7a16 100644
--- a/plugins/aws/fix_plugin_aws/resource/ecr.py
+++ b/plugins/aws/fix_plugin_aws/resource/ecr.py
@@ -75,10 +75,11 @@ def fetch_lifecycle_policy(repository: AwsEcrRepository) -> None:
                 if policy := builder.client.get(
                     service_name,
                     "get-lifecycle-policy",
+                    "lifecyclePolicyText",
                     repositoryName=repository.name,
                     expected_errors=["LifecyclePolicyNotFoundException"],
                 ):
-                    repository.lifecycle_policy = sort_json(json.loads(policy["lifecyclePolicyText"]), sort_list=True)
+                    repository.lifecycle_policy = sort_json(json.loads(policy), sort_list=True)  # type: ignore
 
         def collect(visibility: str, spec: AwsApiSpec) -> None:
             try:
diff --git a/plugins/aws/fix_plugin_aws/resource/s3.py b/plugins/aws/fix_plugin_aws/resource/s3.py
index 78970757a8..b10782d1e1 100644
--- a/plugins/aws/fix_plugin_aws/resource/s3.py
+++ b/plugins/aws/fix_plugin_aws/resource/s3.py
@@ -182,6 +182,7 @@ class AwsS3Bucket(AwsResource, BaseBucket):
     bucket_acl: Optional[AwsS3BucketAcl] = field(default=None)
     bucket_logging: Optional[AwsS3Logging] = field(default=None)
     bucket_location: Optional[str] = field(default=None)
+    bucket_lifecycle_policy: Optional[Json] = field(default=None, metadata={"description": "The bucket lifecycle policy."})  # fmt: skip
 
     @classmethod
     def called_collect_apis(cls) -> List[AwsApiSpec]:
@@ -231,6 +232,19 @@ def add_bucket_policy(bck: AwsS3Bucket) -> None:
                 ):
                     bck.bucket_policy = sort_json(json_loads(raw_policy), sort_list=True)  # type: ignore
 
+        def fetch_lifecycle_policy(bck: AwsS3Bucket) -> None:
+            with builder.suppress(f"{service_name}.get-bucket-lifecycle-configuration"):
+                for policy in builder.client.list(
+                    service_name,
+                    "get-bucket-lifecycle-configuration",
+                    "Rules",
+                    Bucket=bck.name,
+                    expected_errors=["NoSuchLifecycleConfiguration"],
+                ):
+                    if not bck.bucket_lifecycle_policy:
+                        bck.bucket_lifecycle_policy = {}
+                    bck.bucket_lifecycle_policy[policy["ID"]] = policy
+
         def add_bucket_versioning(bck: AwsS3Bucket) -> None:
             with builder.suppress(f"{service_name}.get-bucket-versioning"):
                 if raw_versioning := builder.client.get(
@@ -310,6 +324,7 @@ def add_bucket_location(bck: AwsS3Bucket) -> None:
             builder.submit_work(service_name, add_public_access, bucket)
             builder.submit_work(service_name, add_acls, bucket)
             builder.submit_work(service_name, add_bucket_logging, bucket)
+            builder.submit_work(service_name, fetch_lifecycle_policy, bucket)
 
     def _set_tags(self, client: AwsClient, tags: Dict[str, str]) -> bool:
         tag_set = [{"Key": k, "Value": v} for k, v in tags.items()]
diff --git a/plugins/aws/test/resources/s3_test.py b/plugins/aws/test/resources/s3_test.py
index b4a8858431..645a0f4c47 100644
--- a/plugins/aws/test/resources/s3_test.py
+++ b/plugins/aws/test/resources/s3_test.py
@@ -7,7 +7,7 @@
 
 
 def test_buckets() -> None:
-    first, builder = round_trip_for(AwsS3Bucket)
+    first, builder = round_trip_for(AwsS3Bucket, "bucket_lifecycle_policy")
     assert len(builder.resources_of(AwsS3Bucket)) == 4
     assert len(first.bucket_encryption_rules or []) == 1
     assert first.arn == "arn:aws:s3:::bucket-1"
@@ -25,7 +25,7 @@ def test_s3_account_settings() -> None:
 
 
 def test_tagging() -> None:
-    bucket, _ = round_trip_for(AwsS3Bucket)
+    bucket, _ = round_trip_for(AwsS3Bucket, "bucket_lifecycle_policy")
 
     def validate_update_args(**kwargs: Any) -> Any:
         if kwargs["action"] == "get-bucket-tagging":
@@ -53,7 +53,7 @@ def validate_delete_args(**kwargs: Any) -> Any:
 
 
 def test_deletion() -> None:
-    bucket, _ = round_trip_for(AwsS3Bucket)
+    bucket, _ = round_trip_for(AwsS3Bucket, "bucket_lifecycle_policy")
 
     def validate_delete_args(aws_service: str, fn: Callable[[Any], None]) -> Any:
         assert aws_service == "s3"