From aa60dd25298b59f6859b85eb9daf605c2660778e Mon Sep 17 00:00:00 2001
From: ewandor <ewandor@dorfsvald.net>
Date: Sat, 2 Dec 2023 01:50:01 +0100
Subject: [PATCH] Adding endpoint separation to allow image purging per
 endpoint

---
 main.py | 73 +++++++++++++++++++++++++++++++++++++++------------------
 1 file changed, 50 insertions(+), 23 deletions(-)

diff --git a/main.py b/main.py
index 06111a4..b401121 100644
--- a/main.py
+++ b/main.py
@@ -25,8 +25,13 @@ class PortainerClient:
             "X-API-Key": self.access_token
         }
 
-    def list_stacks(self):
-        response = requests.get(f"{self.api_url}/stacks", headers=self._get_headers())
+    def list_endpoints(self):
+        response = requests.get(f"{self.api_url}/endpoints", headers=self._get_headers())
+        return response.json()
+
+    def list_stacks(self, endpoint_id=None):
+        filters = f'?filters={{"EndpointID":{endpoint_id}}}' if endpoint_id else ""
+        response = requests.get(f"{self.api_url}/stacks{filters}", headers=self._get_headers())
         return response.json()
 
     def get_stack(self, stack_id):
@@ -75,6 +80,12 @@ class PortainerClient:
 
         return f"Downloaded newer image for {image}" in status
 
+    def delete_image(self, endpoint, image):
+        response = requests.delete(
+            f"{self.api_url}/endpoints/{endpoint}/docker/images/{image}?force=false",
+            headers=self._get_headers()
+        )
+
 
 def update_repositories():
     for repo_path in config.git_repositories:
@@ -83,7 +94,19 @@ def update_repositories():
             exit(ret)
 
 
-def update_local_stack(stack):
+def update_stack_images(endpoint, stack_name):
+    containers = client.list_stack_containers(endpoint, stack_name)
+
+    updated_images = []
+    for c in containers:
+        if client.update_image(endpoint, c['Image']):
+            print(f"{c['Image']} was updated.")
+            updated_images.append(c['ImageID'])
+
+    return updated_images
+
+
+def upgrade_local_stack(stack):
     project_name = stack['Name']
 
     print(f"Updating stack: {project_name}")
@@ -96,35 +119,39 @@ def update_local_stack(stack):
     print(f"Stack {project_name} updated!\n")
 
 
-if __name__ == '__main__':
-    update_repositories()
-
-    client = PortainerClient(portainer_url=config.portainer_url, access_token=config.access_token)
-
-    for stack in client.list_stacks():
+def upgrade_endpoint_stack(client, endpoint):
+    endpoint_id = endpoint['Id']
+    endpoint_updated_images = []
+    for stack in client.list_stacks(endpoint_id):
         if stack['Status'] == 1:
-            endpoint = stack['EndpointId']
             stack_name = stack['Name']
 
             print(f"Checking {stack_name} for updates.")
 
-            containers = client.list_stack_containers(endpoint, stack_name)
-
-            an_image_was_updated = False
-            for c in containers:
-                this_image_was_updated = client.update_image(endpoint, c['Image'])
-                if this_image_was_updated:
-                    print(f"{c['Image']} was updated.")
-                    an_image_was_updated = True
-
-            if an_image_was_updated:
+            updated_images = update_stack_images(endpoint_id, stack_name)
+            if updated_images:
+                endpoint_updated_images += updated_images
                 print(f"{stack_name} was updated. Restarting...")
-                if endpoint == 2 and stack_name in ('portainer', 'traefik'):
-                    update_local_stack(stack)
+                if endpoint_id == 2 and stack_name in ('portainer', 'traefik'):
+                    upgrade_local_stack(stack)
                 else:
                     client.upgrade_stack(stack)
                 print(f"{stack_name} restarted.")
             else:
                 print(f"{stack_name} is up to date. Skipping.")
 
-    os.system('docker image prune -fa')
+    if endpoint_updated_images:
+        print("Deleting obsolete images")
+        for image_id in endpoint_updated_images:
+            client.delete_image(endpoint_id, image_id)
+    else:
+        print("No obsolete image to delete")
+
+
+if __name__ == '__main__':
+    update_repositories()
+
+    client = PortainerClient(portainer_url=config.portainer_url, access_token=config.access_token)
+
+    for e in client.list_endpoints():
+        upgrade_endpoint_stack(client, e)