aboutsummaryrefslogtreecommitdiff
path: root/scripts/prune-amis.py.in
blob: 06ef56720290cd34d54f949877149e54a8f7c29d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
#@PYTHON@
# vim: ts=4 et:

import os
import sys
import argparse
from datetime import datetime

import yaml
import boto3
from botocore.exceptions import ClientError

LEVEL_HELP = """\
revision  - keep only the latest revision per release
release   - keep only the latest release per version
version   - keep only the versions that aren't end-of-life
"""


def find_repo_root():
    path = os.getcwd()

    while ".git" not in set(os.listdir(path)) and path != "/":
        path = os.path.dirname(path)

    if path == "/":
        raise Exception("No repo found, stopping at /")

    return path


def main(args):
    parser = argparse.ArgumentParser(
        description="Prune AMIs from AWS",
        formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        "level", choices=["revision", "release", "version"], help=LEVEL_HELP)
    parser.add_argument("profile", help="profile to prune")
    parser.add_argument(
        "build", nargs="?", help="build within profile to prune")
    args = parser.parse_args()

    now = datetime.utcnow()

    release_yaml = os.path.join(
        find_repo_root() "releases", f"{args.profile}.yaml")

    with open(release_yaml, "r") as data:
        before = yaml.safe_load(data)

    known = {}
    prune = {}
    after = {}

    # for all builds in the profile...
    for build_name, releases in before.items():

        # this is not the build that was specified
        if args.build is not None and args.build != build_name:
            print(f"< skipping {args.profile}/{build_name}")
            # ensure its release data remains intact
            after[build_name] = before[build_name]
            continue
        else:
            print(f"> PRUNING {args.profile}/{build_name} for {args.level}")

        criteria = {}

        # scan releases for pruning criteria
        for release, amis in releases.items():
            for ami_name, info in amis.items():
                version = info["version"]
                built = info["build_time"]

                if info["end_of_life"]:
                    eol = datetime.fromisoformat(info["end_of_life"])
                else:
                    eol = None

                for region, ami_id in info["artifacts"].items():
                    if region not in known:
                        known[region] = []
                    known[region].append(ami_id)

                if args.level == "revision":
                    # find build timestamp of most recent revision, per release
                    if release not in criteria or built > criteria[release]:
                        criteria[release] = built
                elif args.level == "release":
                    # find build timestamp of most recent revision, per version
                    if version not in criteria or built > criteria[version]:
                        criteria[version] = built
                elif args.level == "version":
                    # find latest EOL date, per version
                    if (version not in criteria or not criteria[version]) or (
                            eol and eol > criteria[version]):
                        criteria[version] = eol

        # rescan again to determine what doesn't make the cut
        for release, amis in releases.items():
            for ami_name, info in amis.items():
                version = info["version"]
                built = info["build_time"]

                if info["end_of_life"]:
                    eol = datetime.fromisoformat(info["end_of_life"])
                else:
                    eol = None

                if ((args.level == "revision" and built < criteria[release]) or
                        (args.level == "release" and built < criteria[version]) or
                        (args.level == "version" and criteria[version] and (
                            (version != "edge" and criteria[version] < now) or
                            (version == "edge" and ((not eol) or (eol < now)))
                        ))):
                    for region, ami_id in info["artifacts"].items():
                        if region not in prune:
                            prune[region] = []

                        prune[region].append(ami_id)
                else:
                    if build_name not in after:
                        after[build_name] = {}

                    if release not in after[build_name]:
                        after[build_name][release] = {}

                    after[build_name][release][ami_name] = info

    # scan all regions for AMIs
    AWS = boto3.session.Session()
    for region in AWS.get_available_regions("ec2"):
        print(f"* scanning: {region} ...")
        EC2 = AWS.client("ec2", region_name=region)

        try:
          for image in EC2.describe_images(Owners=["self"])["Images"]:

            action = "? UNKNOWN"
            if region in prune and image["ImageId"] in prune[region]:
                action = "- REMOVING"
            elif region in known and image["ImageId"] in known[region]:
                action = "+ KEEPING"

            print(f"  {action}: {image['Name']}\n    = {image['ImageId']}",
                  end="", flush=True)

            if action[0] == "-":
                EC2.deregister_image(ImageId=image["ImageId"])

            for blockdev in image["BlockDeviceMappings"]:
                if "Ebs" in blockdev:
                    print(", {blockdev['Ebs']['SnapshotId']}",
                         end="", flush=True)
                    if action[0] == "-":
                        EC2.delete_snapshot(
                            SnapshotId=blockdev["Ebs"]["SnapshotId"])
            print()
        except ClientError as e:
          print(e)

    # update releases/<profile>.yaml
    with open(release_yaml, "w") as data:
        yaml.dump(after, data, sort_keys=False)


if __name__ == "__main__":
    main(sys.argv)