From 211caada8a0ef03a256a2918ddea72a0f620ec58 Mon Sep 17 00:00:00 2001
From: Pim Kunis <pim@kunis.nl>
Date: Sun, 30 Jun 2024 14:35:47 +0200
Subject: [PATCH] Add script to prefetch Docker images

---
 container-images.nix            |  2 +
 flake-parts/scripts/default.nix | 73 +++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+)
 create mode 100644 container-images.nix

diff --git a/container-images.nix b/container-images.nix
new file mode 100644
index 0000000..037defa
--- /dev/null
+++ b/container-images.nix
@@ -0,0 +1,2 @@
+{ cyberchef = { cyberchef = { finalImageName = "mpepping/cyberchef"; finalImageTag = "latest"; imageDigest = "sha256:5044c72dd8070fb6e0595e720fc4440bf6168493b2433db06a1c966406398ba2"; imageName = "mpepping/cyberchef"; sha256 = "177yjfbz0ijc8lfqfr50fhqqmjk72373c0igyrxv3wwg0pyrgpv4"; }; }; }
+
diff --git a/flake-parts/scripts/default.nix b/flake-parts/scripts/default.nix
index 1148cb7..451e575 100644
--- a/flake-parts/scripts/default.nix
+++ b/flake-parts/scripts/default.nix
@@ -25,4 +25,77 @@ in
     runtimeInputs = with pkgs; [ openssl coreutils openssh yq ];
     scriptPath = ./gen-k3s-cert.sh;
   };
+
+  packages.prefetch-container-images =
+    let
+      images = {
+        cyberchef = {
+          cyberchef = {
+            image-name = "mpepping/cyberchef";
+            image-tag = "latest";
+          };
+        };
+      };
+      imagesJSON = builtins.toFile "images.json" (builtins.toJSON images);
+    in
+    pkgs.writers.writePython3Bin "prefetch-container-images"
+      { } ''
+      import json
+      import subprocess
+      import tempfile
+      import sys
+      from collections import defaultdict
+
+      prefetch_docker_cmd = "${pkgs.lib.getExe pkgs.nix-prefetch-docker}"  # noqa: E501
+      nix_cmd = "${pkgs.lib.getExe pkgs.nix}"  # noqa: E501
+      images_file_name = "${imagesJSON}"
+
+      results = defaultdict(lambda: defaultdict(dict))
+
+      with open(images_file_name, 'r') as file:
+          data = json.load(file)
+
+      for project_name, images in data.items():
+          print(f"Prefetching images for project {project_name}", file=sys.stderr)
+
+          for image_name, image in images.items():
+              name = image["image-name"]
+              tag = image["image-tag"]
+
+              print(f"Prefetching image {name}:{tag}", file=sys.stderr)
+
+              prefetch_args = [
+                  prefetch_docker_cmd,
+                  "--os", "linux",
+                  "--arch", "amd64",
+                  "--image-name", name,
+                  "--image-tag", tag,
+                  "--json",
+                  "--quiet"
+              ]
+              result = subprocess.run(prefetch_args,
+                                      check=True,
+                                      capture_output=True,
+                                      text=True)
+
+              prefetch_data = json.loads(result.stdout)
+              results[project_name][image_name] = prefetch_data
+
+      with tempfile.NamedTemporaryFile(mode='w+', suffix='.json') as temp_file:
+          json.dump(results, temp_file, indent=4)
+          temp_file.flush()
+
+          to_nix_args = [
+              nix_cmd,
+              "eval",
+              "--impure",
+              "--expr", f'builtins.fromJSON (builtins.readFile {temp_file.name})'
+          ]
+          result = subprocess.run(to_nix_args,
+                                  check=True,
+                                  capture_output=True,
+                                  text=True)
+
+      print(result.stdout)
+    '';
 })