nixos-zfs-ec2-ami/manage-amis.ts

#!/usr/bin/env node
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
import { spawnSync } from "node:child_process";
import { fileURLToPath } from "node:url";

const sourceScriptDir = path.dirname(fileURLToPath(import.meta.url));
const repoRoot = process.env.QUIXOS_REPO_ROOT ?? path.resolve(sourceScriptDir, "../..");
const amiBuilderDir = process.env.QUIXOS_AMI_BUILDER_DIR ?? path.resolve(repoRoot, "tofu/ami-builder");
const centralDir = path.resolve(amiBuilderDir, "../central");

const usage = `Usage: nix run ./tofu/ami-builder#ami-artifacts -- <list|delete|publish-ssm> [options]

Commands:
  list                         List published AMI artifacts and deployment references
  delete --ami-id <ami>        Deregister an AMI and delete its root and /nix snapshots
  publish-ssm --ami-id <ami>   Publish an existing AMI into the SSM artifact registry

Options:
  --ami-id <ami>                  AMI to operate on
  --aws-region <region>          AWS region (default: AWS_REGION or us-west-2)
  --name <name>                  AMI project tag (default: nixos-zfs-ec2)
  --architecture <arm64|x86_64>  Filter by architecture
  --runtime-state-bucket <name>  S3 bucket with qx-deploy registry objects
  --force                        Delete even when active registry deployments reference the AMI
  -h, --help                     Show this help
`;

const run = (command, args, options = {}) => {
  const result = spawnSync(command, args, {
    cwd: options.cwd ?? repoRoot,
    env: options.env ?? process.env,
    encoding: "utf8",
    stdio: options.capture ? ["ignore", "pipe", "pipe"] : "inherit",
  });
  if (result.error) {
    throw result.error;
  }
  if (!options.allowFailure && result.status !== 0) {
    const suffix = options.capture && result.stderr ? `\n${result.stderr.trimEnd()}` : "";
    throw new Error(`Command failed: ${command} ${args.join(" ")}${suffix}`);
  }
  return result;
};

const capture = (command, args, options = {}) =>
  run(command, args, { ...options, capture: true }).stdout.trim();

const captureMaybe = (command, args, options = {}) =>
  run(command, args, { ...options, capture: true, allowFailure: true });

const requireCommand = (command) => {
  const result = captureMaybe(command, ["--version"]);
  if (result.status !== 0) {
    throw new Error(`Missing required command: ${command}`);
  }
};

const parseArgs = (argv) => {
  const command = argv[0] ?? "";
  if (!command || command === "-h" || command === "--help") {
    console.error(usage.trimEnd());
    process.exit(command ? 0 : 1);
  }
  if (command !== "list" && command !== "delete" && command !== "publish-ssm") {
    throw new Error(`Unknown command: ${command}\n\n${usage}`);
  }

  const state = {
    command,
    amiId: "",
    awsRegion: process.env.AWS_REGION ?? "us-west-2",
    name: "nixos-zfs-ec2",
    architecture: "",
    runtimeStateBucket: process.env.QUIXOS_RUNTIME_STATE_BUCKET ?? "",
    force: false,
  };

  const takeValue = (args, index, flag) => {
    const value = args[index + 1];
    if (!value) {
      throw new Error(`Missing value for ${flag}`);
    }
    return value;
  };

  for (let i = 1; i < argv.length; i += 1) {
    const arg = argv[i];
    switch (arg) {
      case "--ami-id":
        state.amiId = takeValue(argv, i, arg);
        i += 1;
        break;
      case "--aws-region":
        state.awsRegion = takeValue(argv, i, arg);
        i += 1;
        break;
      case "--name":
        state.name = takeValue(argv, i, arg);
        i += 1;
        break;
      case "--architecture":
        state.architecture = takeValue(argv, i, arg);
        i += 1;
        break;
      case "--runtime-state-bucket":
        state.runtimeStateBucket = takeValue(argv, i, arg);
        i += 1;
        break;
      case "--force":
        state.force = true;
        break;
      case "-h":
      case "--help":
        console.error(usage.trimEnd());
        process.exit(0);
      default:
        throw new Error(`Unknown argument: ${arg}\n\n${usage}`);
    }
  }

  return state;
};

const resolveRuntimeStateBucket = (state) => {
  if (state.runtimeStateBucket) {
    return;
  }
  if (!fs.existsSync(centralDir)) {
    return;
  }
  const result = captureMaybe("tofu", ["-chdir=" + centralDir, "output", "-raw", "runtime_state_bucket_name"]);
  if (result.status === 0 && /^[a-z0-9][a-z0-9.-]{1,61}[a-z0-9]$/.test(result.stdout.trim())) {
    state.runtimeStateBucket = result.stdout.trim();
  }
};

const awsJson = (args) => {
  const text = capture("aws", [...args, "--output", "json"]);
  return text ? JSON.parse(text) : null;
};

const amiTags = (image) => Object.fromEntries((image.Tags ?? []).map((tag) => [tag.Key, tag.Value]));

const listImages = (state) => {
  const filters = [
    "Name=state,Values=available",
    `Name=tag:Project,Values=${state.name}`,
    "Name=tag:Role,Values=published-ami",
  ];
  if (state.architecture) {
    filters.push(`Name=tag:Architecture,Values=${state.architecture}`);
  }

  const result = awsJson([
    "ec2",
    "describe-images",
    "--region",
    state.awsRegion,
    "--owners",
    "self",
    "--filters",
    ...filters,
  ]);

  return (result?.Images ?? [])
    .sort((a, b) => String(b.CreationDate ?? "").localeCompare(String(a.CreationDate ?? "")));
};

const deploymentRegistry = (state) => {
  if (!state.runtimeStateBucket) {
    return [];
  }

  const listed = captureMaybe("aws", [
    "s3api",
    "list-objects-v2",
    "--region",
    state.awsRegion,
    "--bucket",
    state.runtimeStateBucket,
    "--prefix",
    "registry/deployments/",
    "--query",
    "Contents[].Key",
    "--output",
    "text",
  ]);
  if (listed.status !== 0 || !listed.stdout.trim() || listed.stdout.trim() === "None") {
    return [];
  }

  const records = [];
  const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "quixos-ami-registry."));
  try {
    for (const key of listed.stdout.trim().split(/\s+/)) {
      if (!key.endsWith(".json")) {
        continue;
      }
      const result = captureMaybe("aws", ["s3", "cp", `s3://${state.runtimeStateBucket}/${key}`, "-"]);
      if (result.status === 0 && result.stdout.trim()) {
        records.push(JSON.parse(result.stdout));
      }
    }
  } finally {
    fs.rmSync(tmpDir, { recursive: true, force: true });
  }
  return records;
};

const deploymentReferenceLabel = (record) =>
  `${record.id}${record.status ? `:${record.status}` : ""}`;

const isDestroyedDeployment = (record) => record.status === "destroyed";

const listArtifacts = (state) => {
  const registry = deploymentRegistry(state);
  const references = new Map();
  for (const record of registry) {
    const amiId = record.config?.amiId;
    if (!amiId) {
      continue;
    }
    const entries = references.get(amiId) ?? [];
    entries.push(deploymentReferenceLabel(record));
    references.set(amiId, entries);
  }

  for (const image of listImages(state)) {
    const tags = amiTags(image);
    const refs = references.get(image.ImageId)?.join(",") ?? "-";
    console.log([
      image.ImageId,
      image.Name ?? "",
      image.CreationDate ?? "",
      tags.Architecture ?? image.Architecture ?? "",
      tags.BuilderGitRev ?? "",
      tags.RootSnapshotId ?? "",
      tags.NixSnapshotId ?? "",
      refs,
    ].join("\t"));
  }
};

const deleteArtifact = (state) => {
  if (!state.amiId) {
    throw new Error("Missing --ami-id");
  }

  const registry = deploymentRegistry(state);
  const activeRefs = registry
    .filter((record) => record.config?.amiId === state.amiId)
    .filter((record) => !isDestroyedDeployment(record))
    .map(deploymentReferenceLabel);
  if (activeRefs.length > 0 && !state.force) {
    throw new Error(`Refusing to delete ${state.amiId}; referenced by active deployments: ${activeRefs.join(", ")}\nDestroy those deployments first, or pass --force if they are intentionally orphaned.`);
  }

  const result = awsJson([
    "ec2",
    "describe-images",
    "--region",
    state.awsRegion,
    "--owners",
    "self",
    "--image-ids",
    state.amiId,
  ]);
  const image = result?.Images?.[0];
  if (!image) {
    throw new Error(`AMI not found: ${state.amiId}`);
  }
  const tags = amiTags(image);
  const rootSnapshot = tags.RootSnapshotId ?? "";
  const nixSnapshot = tags.NixSnapshotId ?? "";

  console.error(`Deregistering AMI: ${state.amiId}`);
  run("aws", ["ec2", "deregister-image", "--region", state.awsRegion, "--image-id", state.amiId]);

  if (rootSnapshot) {
    console.error(`Deleting root snapshot: ${rootSnapshot}`);
    run("aws", ["ec2", "delete-snapshot", "--region", state.awsRegion, "--snapshot-id", rootSnapshot]);
  }
  if (nixSnapshot) {
    console.error(`Deleting /nix snapshot: ${nixSnapshot}`);
    run("aws", ["ec2", "delete-snapshot", "--region", state.awsRegion, "--snapshot-id", nixSnapshot]);
  }
};

const getImage = (state) => {
  const result = awsJson([
    "ec2",
    "describe-images",
    "--region",
    state.awsRegion,
    "--owners",
    "self",
    "--image-ids",
    state.amiId,
  ]);
  const image = result?.Images?.[0];
  if (!image) {
    throw new Error(`AMI not found: ${state.amiId}`);
  }
  return image;
};

const putSsmParameter = (state, name, value) => {
  run("aws", [
    "ssm",
    "put-parameter",
    "--region",
    state.awsRegion,
    "--name",
    name,
    "--type",
    "String",
    "--value",
    value,
    "--overwrite",
  ]);
};

const publishSsmArtifact = (state) => {
  if (!state.amiId) {
    throw new Error("Missing --ami-id");
  }

  const image = getImage(state);
  const tags = amiTags(image);
  const architecture = tags.Architecture ?? image.Architecture ?? "";
  const builderGitRev = tags.BuilderGitRev ?? "";
  const inputHash = tags.InputHash ?? "";
  if (!architecture) {
    throw new Error(`AMI ${state.amiId} is missing Architecture metadata`);
  }
  if (!builderGitRev) {
    throw new Error(`AMI ${state.amiId} is missing BuilderGitRev tag`);
  }
  if (!inputHash) {
    throw new Error(`AMI ${state.amiId} is missing InputHash tag`);
  }

  const prefix = `/quixos/amis/${state.name}/${architecture}`;
  const parameters = [
    `${prefix}/recommended`,
    `${prefix}/by-builder-rev/${builderGitRev}`,
    `${prefix}/by-input-hash/${inputHash}`,
  ];
  for (const parameter of parameters) {
    console.error(`Writing ${parameter} = ${state.amiId}`);
    putSsmParameter(state, parameter, state.amiId);
  }
};

const main = () => {
  const state = parseArgs(process.argv.slice(2));
  requireCommand("aws");
  resolveRuntimeStateBucket(state);

  if (state.command === "list") {
    listArtifacts(state);
    return;
  }
  if (state.command === "publish-ssm") {
    publishSsmArtifact(state);
    return;
  }
  deleteArtifact(state);
};

try {
  main();
} catch (error) {
  console.error(error instanceof Error ? error.message : String(error));
  process.exit(1);
}