From f98b92073da22d5897c019bcfe02036e74709e13 Mon Sep 17 00:00:00 2001
From: Seth Foster <fosterbseth@gmail.com>
Date: Wed, 14 Apr 2021 14:46:51 -0400
Subject: [PATCH] Add cleanup_images system job template

- Removes podman images on the system that are not assigned to an
execution environment
---
 .../management/commands/cleanup_images.py     | 79 +++++++++++++++++++
 awx/main/migrations/0136_cleanup_ee_images.py | 32 ++++++++
 2 files changed, 111 insertions(+)
 create mode 100644 awx/main/management/commands/cleanup_images.py
 create mode 100644 awx/main/migrations/0136_cleanup_ee_images.py

diff --git a/awx/main/management/commands/cleanup_images.py b/awx/main/management/commands/cleanup_images.py
new file mode 100644
index 0000000000..9d61d353a3
--- /dev/null
+++ b/awx/main/management/commands/cleanup_images.py
@@ -0,0 +1,79 @@
+# Copyright (c) 2015 Ansible, Inc.
+# All Rights Reserved.
+
+# Python
+import subprocess
+import logging
+import json
+
+# Django
+from django.core.management.base import BaseCommand, CommandError
+from django.conf import settings
+
+# AWX
+from awx.main.models import ExecutionEnvironment
+
+
+class Command(BaseCommand):
+    """
+    Management command to cleanup unused execution environment images.
+    """
+
+    help = 'Remove unused execution environment images'
+
+    def init_logging(self):
+        log_levels = dict(enumerate([logging.ERROR, logging.INFO, logging.DEBUG, 0]))
+        self.logger = logging.getLogger('awx.main.commands.cleanup_sessions')
+        self.logger.setLevel(log_levels.get(self.verbosity, 0))
+        handler = logging.StreamHandler()
+        handler.setFormatter(logging.Formatter('%(message)s'))
+        self.logger.addHandler(handler)
+        self.logger.propagate = False
+
+    def add_arguments(self, parser):
+        parser.add_argument('--dry-run', dest='dry_run', action='store_true', default=False, help='Dry run mode (show items that would ' 'be removed)')
+
+    def delete_images(self, images_json):
+        for e in images_json:
+            if 'Names' in e:
+                image_names = e['Names']
+            else:
+                image_names = [e["Id"]]
+            image_size = e['Size'] / 1e6
+            for i in image_names:
+                if i not in self.images_in_use and i not in self.deleted:
+                    self.deleted.append(i)
+                    self.logger.info(f"{self.delete_prefix} {i}: {image_size:.0f} MB")
+                    if not self.dry_run:
+                        subprocess.run(['podman', 'rmi', i, '-f'], stdout=subprocess.DEVNULL)
+
+    def cleanup_images(self):
+        self.images_in_use = [ee.image for ee in ExecutionEnvironment.objects.all()]
+        self.logger.info(f"Execution environment images in use: {self.images_in_use}")
+        self.deleted = []
+        # find and remove unused images
+        images_system = subprocess.run("podman images -a --format json".split(" "), capture_output=True)
+        if len(images_system.stdout) > 0:
+            images_system = json.loads(images_system.stdout)
+
+            self.delete_images(images_system)
+        # find and remove dangling images
+        images_system = subprocess.run('podman images -a --filter "dangling=true" --format json'.split(" "), capture_output=True)
+        if len(images_system.stdout) > 0:
+            images_system = json.loads(images_system.stdout)
+            self.delete_images(images_system)
+        if not self.deleted:
+            self.logger.info("Did not find images to remove")
+
+    def handle(self, *args, **options):
+        self.verbosity = int(options.get('verbosity', 1))
+        self.init_logging()
+        self.dry_run = bool(options.get('dry_run', False))
+        if self.dry_run:
+            self.delete_prefix = "Would delete"
+            self.logger.info("Dry run enabled, images will not be deleted")
+        else:
+            self.delete_prefix = "Deleting"
+        if settings.IS_K8S:
+            raise CommandError("Cannot run cleanup tool on k8s installations")
+        self.cleanup_images()
diff --git a/awx/main/migrations/0136_cleanup_ee_images.py b/awx/main/migrations/0136_cleanup_ee_images.py
new file mode 100644
index 0000000000..86b81a0a35
--- /dev/null
+++ b/awx/main/migrations/0136_cleanup_ee_images.py
@@ -0,0 +1,32 @@
+# Generated by Django 2.2.16 on 2021-04-14 16:21
+
+from django.db import migrations
+from django.utils.timezone import now
+
+
+def create_cleanup_ee_images(apps, schema_editor):
+    SystemJobTemplate = apps.get_model('main', 'SystemJobTemplate')
+    ContentType = apps.get_model('contenttypes', 'ContentType')
+    sjt_ct = ContentType.objects.get_for_model(SystemJobTemplate)
+    now_dt = now()
+    sjt, created = SystemJobTemplate.objects.get_or_create(
+        job_type='cleanup_images',
+        defaults=dict(
+            name='Cleanup Execution Environment Images',
+            description='Remove unused execution environment images',
+            created=now_dt,
+            modified=now_dt,
+            polymorphic_ctype=sjt_ct,
+        ),
+    )
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('main', '0135_schedule_sort_fallback_to_id'),
+    ]
+
+    operations = [
+        migrations.RunPython(create_cleanup_ee_images),
+    ]