disko

#!/usr/bin/env bash
set -euo pipefail

readonly libexec_dir="${0%/*}"

# a file with the disko config
declare disko_config

# mount was chosen as the default mode because it's less destructive
mode=mount
nix_args=()
skip_destroy_safety_check=false

# DISKO_VERSION is set by the wrapper in package.nix
DISKO_VERSION="${DISKO_VERSION:="unknown! This is a bug, please report it!"}"
onlyPrintVersion=false

showUsage() {
  cat <<USAGE
Usage: $0 [options] disk-config.nix
or $0 [options] --flake github:somebody/somewhere#disk-config

With flakes, disk-config is discovered first under the .diskoConfigurations top level attribute
or else from the disko module of a NixOS configuration of that name under .nixosConfigurations.

Options:

* -m, --mode mode
  set the mode, either distroy, format, mount, format,mount or destroy,format,mount
    destroy: unmount filesystems and destroy partition tables of the selected disks
    format: create partition tables, zpools, lvms, raids and filesystems if they don't exist yet
    mount: mount the partitions at the specified root-mountpoint
    format,mount: run format and mount in sequence
    destroy,format,mount: run all three modes in sequence. Previously known as --mode disko
* -f, --flake uri
  fetch the disko config relative to this flake's root
* --arg name value
  pass value to nix-build. can be used to set disk-names for example
* --argstr name value
  pass value to nix-build as string
* --root-mountpoint /some/other/mnt
  where to mount the device tree (default: /mnt)
* --dry-run
  just show the path to the script instead of running it
* --no-deps
  avoid adding another dependency closure to an in-memory installer
    requires all necessary dependencies to be available in the environment
* --yes-wipe-all-disks
  skip the safety check for destroying partitions, useful for automation
* --debug
  run with set -x
* --help
  show this help
USAGE
}

abort() {
  echo "aborted: $*" >&2
  exit 1
}

## Main ##

[[ $# -eq 0 ]] && {
  showUsage
  exit 1
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --debug)
      set -x
      ;;
    -m | --mode)
      mode=$2
      shift
      ;;
    -f | --flake)
      flake=$2
      shift
      ;;
    --argstr | --arg)
      nix_args+=("$1" "$2" "$3")
      shift
      shift
      ;;
    -h | --help)
      showUsage
      exit 0
      ;;
    --dry-run)
      dry_run=y
      ;;
    --root-mountpoint)
      nix_args+=(--argstr rootMountPoint "$2")
      shift
      ;;
    --no-deps)
      nix_args+=(--arg noDeps true)
      ;;
    --yes-wipe-all-disks)
      skip_destroy_safety_check=true
      ;;
    --show-trace)
      nix_args+=("$1")
      ;;
    --version)
      onlyPrintVersion=true
      ;;
    *)
      if [ -z ${disko_config+x} ]; then
        disko_config=$1
      else
        showUsage
        exit 1
      fi
      ;;
  esac
  shift
done

if [[ "$onlyPrintVersion" = true ]]; then
  echo "$DISKO_VERSION"
  exit 0
fi
# Always print version information to help with debugging
echo "disko version $DISKO_VERSION"

nixBuild() {
  if command -v nom-build > /dev/null; then
    nom-build "$@"
  else
    nix-build "$@"
  fi
}

if ! {
  # Base modes
  [[ $mode = "destroy" ]] || [[ $mode = "format" ]] || [[ $mode = "mount" ]] ||
  # Combined modes
  [[ $mode = "format,mount" ]] ||
  [[ $mode = "destroy,format,mount" ]] || # Replaces --mode disko
  # Legacy modes, will be removed in next major version
  [[ $mode = "disko" ]] || [[ $mode = "create" ]] || [[ $mode = "zap_create_mount" ]] ;
}; then
  abort 'mode must be one of "destroy", "format", "mount", "destroy,format,mount" or "format,mount"'
fi

if [[ -n "${flake+x}" ]]; then
  if [[ $flake =~ ^(.*)\#([^\#\"]*)$ ]]; then
   flake="${BASH_REMATCH[1]}"
   flakeAttr="${BASH_REMATCH[2]}"
  fi
  if [[ -z "${flakeAttr-}" ]]; then
    echo "Please specify the name of the NixOS configuration to be installed, as a URI fragment in the flake-uri."
    echo "For example, to use the output diskoConfigurations.foo from the flake.nix, append \"#foo\" to the flake-uri."
    exit 1
  fi
  if [[ -e "$flake" ]]; then
    flake="$(realpath "$flake")"
  fi
  nix_args+=("--arg" "flake" "\"$flake\"")
  nix_args+=("--argstr" "flakeAttr" "$flakeAttr")
  nix_args+=(--extra-experimental-features flakes)
elif [[ -n "${disko_config+x}" ]] && [[ -e "$disko_config" ]]; then
  nix_args+=("--arg" "diskoFile" "$(realpath "$disko_config")")
else
  abort "disko config must be an existing file or flake must be set"
fi

# The "--impure" is still pure, as the path is within the nix store.
script=$(nixBuild "${libexec_dir}"/cli.nix \
  --no-out-link \
  --impure \
  --argstr mode "$mode" \
  "${nix_args[@]}"
)

command=("$(echo "$script"/bin/*)")
if [[ $mode = "destroy,format,mount" && $skip_destroy_safety_check = true ]]; then
  command+=("--yes-wipe-all-disks")
fi

if [[ -n "${dry_run+x}" ]]; then
  echo "${command[@]}"
else
  exec "${command[@]}"
fi