nixos-servers/scripts/default.nix

{ self, nixpkgs, flake-utils, ... }: flake-utils.lib.eachDefaultSystem (system:
let
  pkgs = nixpkgs.legacyPackages.${system};
  createScript = { name, runtimeInputs, scriptPath, extraWrapperFlags ? "", ... }:
    let
      script = (pkgs.writeScriptBin name (builtins.readFile scriptPath)).overrideAttrs (old: {
        buildCommand = "${old.buildCommand}\n patchShebangs $out";
      });
    in
    pkgs.symlinkJoin {
      inherit name;
      paths = [ script ] ++ runtimeInputs;
      buildInputs = [ pkgs.makeWrapper ];
      postBuild = "wrapProgram $out/bin/${name} --set PATH $out/bin ${extraWrapperFlags}";
    };
in
{
  packages.bootstrap = createScript {
    name = "bootstrap";
    runtimeInputs = with pkgs; [ sops coreutils nixos-anywhere ];
    scriptPath = ./bootstrap.sh;
  };

  packages.gen-k3s-cert = createScript {
    name = "create-k3s-cert";
    runtimeInputs = with pkgs; [ openssl coreutils openssh yq ];
    scriptPath = ./gen-k3s-cert.sh;
  };

  packages.prefetch-container-images =
    let
      imagesJSON = builtins.toFile "images.json" (builtins.toJSON self.globals.images);
    in
    pkgs.writers.writePython3Bin "prefetch-container-images.py"
      { } ''
      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 image_name, image_ref in data.items():
          [name, tag] = image_ref.split(":", maxsplit=1)
          print(f"Prefetching image {image_ref}", file=sys.stderr)

          digest = ""
          if "@" in tag:
              [tag, digest] = tag.split("@", maxsplit=1)

          prefetch_args = [
              prefetch_docker_cmd,
              "--os", "linux",
              "--arch", "amd64",
              "--image-name", name,
              "--image-tag", tag,
              "--json",
              "--quiet"
          ]

          if digest:
              prefetch_args.extend(["--image-digest", digest])

          result = subprocess.run(prefetch_args,
                                  check=True,
                                  capture_output=True,
                                  text=True)

          prefetch_data = json.loads(result.stdout)
          results[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)
    '';
})