Adding garbage removal for cloud uploads

Author: gursewak1997

src/remote_prune

@@ -0,0 +1,236 @@
+#!/usr/bin/python3 -u
+
+import argparse
+import subprocess
+import json
+import yaml
+import collections
+import datetime
+import os
+from dateutil.relativedelta import relativedelta
+
+from cosalib.aliyun import remove_aliyun_image
+from cosalib.gcp import remove_gcp_image
+from cosalib.azure import remove_azure_image
+from cosalib.prune import fetch_build_meta, delete_build
+
+from cosalib.aws import deregister_ami, delete_snapshot
+
+Build = collections.namedtuple("Build", ["id", "timestamp", "images", "arches"])
+
+
+def main():
+    parser = argparse.ArgumentParser(prog="coreos-assembler remote-prune")
+    parser.add_argument("--policy", dest="policy", type=str, help="Path to policy.yaml file")
+    parser.add_argument("--dry-run", dest='dry_run', help="Don't actually delete anything",
+                        action='store_true')
+    parser.add_argument("--stream", dest="stream", type=str, help="Fedora stream", required=True)
+    parser.add_argument('--azure-credentials', help='Path to Azure credentials file',
+                    default=os.environ.get("AZURE_CREDENTIALS"))
+    parser.add_argument('--azure-resource-group', help='Resource group',
+                        default=os.environ.get('AZURE_RESOURCE_GROUP'))
+    parser.add_argument("--gcp-json-key", help="GCP Service Account JSON Auth",
+                        default=os.environ.get("GCP_JSON_AUTH"))
+    parser.add_argument("--gcp-project", help="GCP Project name",
+                        default=os.environ.get("GCP_PROJECT_NAME"))
+    parser.add_argument("--aws-credentials",help="AWS Credentials",
+                        default=os.environ.get("GCP_PROJECT_NAME"))
+
+    # subparsers = parser.add_subparsers(dest='cmd', title='subcommands')
+    # subparsers.required = True
+
+    # s3 = subparsers.add_parser('s3', help='Prune s3 buckets')
+    # s3.add_argument("--bucket", help="Bucket name")
+    # s3.add_argument("--prefix", help="Key prefix")
+    # s3.add_argument("--force", help="Wipe s3 key ignoring the errors",
+    #                 action='store_true')
+
+    args = parser.parse_args()
+
+    cloud_config = {
+        'azure': {
+            'credentials': args.azure_credentials,
+            'resource-group': args.azure_resource_group,
+        },
+        'gcp': {
+            'json-key': args.gcp_json_key,
+            'project': args.gcp_project,
+        }
+    }
+
+    policy = open(f"{args.policy}")
+    policy = yaml.safe_load(policy)
+    stream = args.stream
+
+    if stream in policy:
+        # If the build key is set in the policy file, then the cloud-uploads key must
+        # also be present, and the duration of cloud-uploads must be equal or shorter
+        if "build" in policy[stream].keys():
+            cloudUploadsCheck(policy[stream])
+        buildJsonData = getBuildsForStream(stream)
+        # action is basically whatever is needed to be pruned for the respective stream
+        for action in policy[stream]:
+            years = policy[stream][action]
+            #Ref date: Artifacts older than this date will be deleted
+            refDate = datetime.datetime.now() - relativedelta(years=int(years[:-1]))
+            print(f"Reference Date for action {action} is {refDate}")
+
+            for index, build in enumerate(buildJsonData["builds"]):
+                build_id = build["id"]
+                if "policy-cleanup" in build.keys():
+                    if action in build["policy-cleanup"]:
+                        print(f"The {build_id} has already had {action} pruning completed")
+                        break
+                timestamp = build_id.split('.')[1]
+                buildDate = datetime.datetime(int(timestamp[0:4]), int(timestamp[4:6]), int(timestamp[-2:]))
+                if buildDate < refDate:
+                    # For testing pusposes only using one build
+                    if build_id == "35.20210508.91.0":
+                        for arch in build["arches"]:
+                            print(f"Pruning {years} old {action} for {build_id} for {arch} in {stream}")
+                            # buildFetch(args.stream, build_id, arch)
+                            metaJson = open(f"builds/{build_id}/{arch}/meta.json")
+                            metaData = json.load(metaJson)
+
+                            images = {
+                                "amis": metaData.get("amis") or [],
+                                "azure": metaData.get("azure") or [],
+                                "gcp": metaData.get("gcp") or [],
+                            }
+                            currentBuild = Build(
+                                id=build_id,
+                                timestamp=timestamp,
+                                images=images,
+                                arches=arch,
+                            )
+                            match action:
+                                case "cloud-uploads":
+                                    delete_cloud_resource(currentBuild, cloud_config, args.dry_run)
+                                    if not args.dry_run:
+                                        build.setdefault("policy-cleanup", []).append(action)
+                                        buildJsonData["builds"][index] = build
+                                case "build":
+                                    print(f"Will delete everything for {build_id}")
+        # pprint.pprint(buildJsonData['builds'], compact=True)
+        with open(f"builds/builds.json", "w") as json_file:
+            json_file.write(json.dumps(buildJsonData))
+
+def delete_cloud_resource(build, cloud_config, dry_run):
+    print(f"Deleting cloud uploads for {build.id}")
+    errors = []
+    # Unregister AMIs and snapshots
+    for ami in build.images.get("amis", []):
+        region_name = ami.get("name")
+        ami_id = ami.get("hvm")
+        snapshot_id = ami.get("snapshot")
+        if dry_run:
+            print(f"Will delete {ami_id} and {snapshot_id}")
+        else:
+            if ami_id and region_name:
+                try:
+                    deregister_ami(ami_id, region=region_name)
+                except Exception as e:
+                    errors.append(e)
+            if snapshot_id and region_name:
+                try:
+                    delete_snapshot(snapshot_id, region=region_name)
+                except Exception as e:
+                    errors.append(e)
+
+    aliyun = build.images.get("aliyun")
+    if aliyun:
+        region_name = aliyun.get("name")
+        aliyun_id = aliyun.get("hvm")
+        if dry_run:
+            print(f"Will delete {aliyun_id} in {region_name}")
+        else:
+            if region_name and aliyun_id:
+                try:
+                    remove_aliyun_image(aliyun_id, region=region_name)
+                except Exception as e:
+                    errors.append(e)
+
+    azure = build.images.get("azure")
+    if azure:
+        image = azure.get("image")
+        resource_group = cloud_config.get("azure", {}).get("resource-group")
+        credentials = cloud_config.get("azure", {}).get("credentials")
+        if dry_run:
+            print(f"Will delete Azure {image}")
+        else:
+            if image and resource_group and credentials:
+                try:
+                    remove_azure_image(image, resource_group, credentials)
+                except Exception as e:
+                    errors.append(e)
+
+    gcp = build.images.get("gcp")
+    if gcp:
+        gcp_image = gcp.get("image")
+        json_key = cloud_config.get("gcp", {}).get("json-key")
+        project = cloud_config.get("gcp", {}).get("project")
+        if dry_run:
+            print(f"Will delete {gcp_image}")
+        else:
+            if gcp_image and json_key and project:
+                try:
+                    remove_gcp_image(gcp_image, json_key, project)
+                except Exception as e:
+                    errors.append(e)
+
+    if len(errors) != 0:
+        print(f"Found errors when removing build {build.id}:")
+        for e in errors:
+            print(e)
+
+def cloudUploadsCheck(actions):
+    if "cloud-uploads" in actions.keys():
+        cloud_uploads_duration = actions["cloud-uploads"]
+        build_duration = actions["build"]
+        # assumption we are keeping the duration in years
+        assert cloud_uploads_duration < build_duration
+    else:
+        print(
+            "cloud-uploads must be set or be less than builds pruning duration in policy.yaml"
+        )
+
+
+def getBuildsForStream(stream):
+    # buildFetchCmd = 'cosa buildfetch --stream='+ stream + ' --arch=all'
+    # subprocess.check_output(buildFetchCmd.split(' '), shell=True)
+
+    f = open(f"builds/builds.json")
+    buildJsonData = json.load(f)
+    return buildJsonData
+
+
+def buildFetch(stream, build, arch):
+    buildFetchCmd = (
+        "cosa buildfetch --stream=" + stream + " --build=" + build + " --arch=" + arch
+    )
+    subprocess.check_output(buildFetchCmd.split(" "))
+
+#Ignore this
+def awsPrune(amis, dry_run):
+    print(f"The AMIs are {amis}")
+    errors = []
+    for ami in amis:
+        region_name = ami.get("name")
+        ami_id = ami.get("hvm")
+        snapshot_id = ami.get("snapshot")
+        if dry_run:
+            print(f"Would prune {ami_id} and {snapshot_id} in {region_name}")
+        else:
+            if ami_id and region_name:
+                try:
+                    deregister_ami(ami_id, region=region_name)
+                except Exception as e:
+                    errors.append(e)
+            if snapshot_id and region_name:
+                try:
+                    delete_snapshot(snapshot_id, region=region_name)
+                except Exception as e:
+                    errors.append(e)
+
+if __name__ == "__main__":
+    main()