# auto-upgrade-with-pinch: Secure managed NixOS updates
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation, version 3.

{
  config,
  lib,
  pkgs,
  ...
}:
with lib;
let
  local-pkgs = import ../. { inherit pkgs; };
  cfg = config.system.autoUpgradeWithPinch;
  pull-repo-script = pkgs.writeShellScript "pull-repo" ''
    set -eo pipefail

    path=$1
    config=$2

    prop() {
      ${pkgs.jq}/bin/jq -r ".$1" <<< "$config"
    }

    echo Pulling in "$path" >&2

    if [[ ! -e "$path" ]];then
      d=$(mktemp -d)
      ${pkgs.git}/bin/git init "$d"
      ${pkgs.git}/bin/git -C "$d" checkout -b "$(prop localBranch)"
      ${pkgs.git}/bin/git -C "$d" remote add "$(prop remoteName)" "$(prop url)"
      ${pkgs.git}/bin/git -C "$d" branch -u "$(prop remoteBranch)"
      mkdir -p "$(${pkgs.coreutils}/bin/dirname "$path")"
      mv "$d" "$path"
    fi

    cd "$path"

    if [[ "$(${pkgs.git}/bin/git remote get-url "$(prop remoteName)")" != "$(prop url)" ]]; then
      echo Expected git remote "$(prop remoteName)" to point at "$(prop url)" \
        but it points at "$(${pkgs.git}/bin/git remote get-url "$(prop remoteName)")" >&2
      case "$(prop onRemoteURLMismatch)" in
        abort)  exit 1;;
        update) echo Updating it >&2
                ${pkgs.git}/bin/git -C "$d" remote set-url "$(prop remoteName)" "$(prop url)";;
      esac
    fi

    ${pkgs.git}/bin/git fetch "$(prop remoteName)" "$(prop remoteBranch)"

    if [[ "$(${pkgs.git}/bin/git rev-parse --abbrev-ref HEAD)" != "$(prop localBranch)" ]];then
      echo Could not merge because currently-checked-out \
           \""$(${pkgs.git}/bin/git rev-parse --abbrev-ref HEAD)"\" is not \
           \""$(prop localBranch)"\"
      case "$(prop onBranchMismatch)" in
        abort) exit 1;;
        continue) exit 0;;
      esac
    fi

    if [[ "$(prop requireSignature)" == true ]]; then
      ${pkgs.polite-merge}/bin/polite-merge \
        -c gpg.program=${escapeShellArg (local-pkgs.keyed-gpg cfg.signingKeys)} \
        merge --ff-only --verify-signatures
    else
      ${pkgs.polite-merge}/bin/polite-merge merge --ff-only
    fi
  '';

  auto-upgrade-script = pkgs.writeShellScript "auto-upgrade" ''
    ${pkgs.coreutils}/bin/nice -n 17 \
    ${pkgs.util-linux}/bin/ionice -c 3 \
    ${pkgs.util-linux}/bin/flock /run/auto-upgrade-with-pinch ${pkgs.writeShellScript "auto-upgrade-with-lock-held" ''
      set -eo pipefail

      dry_run=false
      pinch_args=()
      if [[ "$1" == --dry-run ]];then
        dry_run=true
        pinch_args=( --dry-run )
      fi

      hydrate() {
        if "$dry_run";then
          echo "Would run: $*"
        else
          "$@"
        fi
      }

      die() {
        echo "$*" >&2
        exit 1
      }

      in_tmpdir() {
        d=$(${pkgs.coreutils}/bin/mktemp -d)
        pushd "$d"
        "$@"
        popd
        ${pkgs.coreutils}/bin/rm -r "$d"
      }

      ${optionalString (cfg.upgradeConfigOwnershipPolicy != "any") (
        ''
          verify_ownership() {
            if [[ "$1" != /* ]];then
              die "Unexpected relative path: $1"
            fi
            if [[ "$1" != / ]];then
              verify_ownership "$(${pkgs.coreutils}/bin/dirname "$1")"
            fi
            if [[ ! -e "$1" ]];then
              die "Could not find upgrade config: $1 does not exist"
            fi
            if [[ -h "$1" ]];then
              verify_ownership "$(
                ${pkgs.coreutils}/bin/realpath --no-symlinks \
                  "$(${pkgs.coreutils}/bin/dirname "$1")/$(${pkgs.coreutils}/bin/readlink "$1")"
              )"
            fi
            perms="$(${pkgs.findutils}/bin/find "$1" -maxdepth 0 -printf "%M")"
            if [[ "$perms" == d*t ]];then
              die "Will not use upgrade config in sticky directory $1"
            fi
            owner=$(${pkgs.findutils}/bin/find "$1" -maxdepth 0 -printf "%u")
            if [[ "$owner" != root ]];then
              die "Will not use upgrade config not owned by root in $1"
            fi
            if [[ "$perms" == l* ]];then
              return 0 # Root-owned symlinks are fine
            fi
            if [[ "$perms" == *w? ]];then
              die "Will not use world-writable upgrade config in $1"
            fi
            ${
              {
                root = ''
                  if [[ "$perms" == *w???? ]];then
                    die "Will not use group-writable upgrade config in $1"
                  fi
                '';
                wheel = ''
                  if [[ "$perms" == *w???? ]];then
                    group=$(${pkgs.findutils}/bin/find "$1" -maxdepth 0 -printf "%g")
                    if [[ "$group" != wheel ]];then
                      die "Will not use non-wheel-group group-writable upgrade config in $1"
                    fi
                  fi
                '';
              }
              ."${cfg.upgradeConfigOwnershipPolicy}"
            }
          }
        ''
        + concatMapStringsSep "\n" (f: "verify_ownership ${escapeShellArg f}") cfg.upgradeConfig
      )}

      config=$(${pkgs.nix}/bin/nix-instantiate --eval --strict --json -A config \
        --arg upgradeConfig ${
          escapeShellArg (
            "[" + lib.concatMapStringsSep " " lib.strings.escapeNixString cfg.upgradeConfig + "]"
          )
        } ${../upgrade-config.nix})

      config_query() {
        ${pkgs.jq}/bin/jq -r "$@" <<< "$config"
      }

      repo_query() {
        config_query --arg path "$1" ".repos[\$ARGS.named.path]$2"
      }

      userenv_query() {
        config_query --arg user "$1" ".userEnvironments[\$ARGS.named.user]$2"
      }

      # Pull updates
      while read path;do
        hydrate /run/wrappers/bin/sudo -u "$(repo_query "$path" .user)" \
          ${pull-repo-script} "$path" "$(repo_query "$path" "")"
      done < <( config_query '.repos | keys []' )

      # Update channels
      config_query '.pinchFiles[]' | ${pkgs.findutils}/bin/xargs --no-run-if-empty --delimiter=\\n ${pkgs.pinch}/bin/pinch update "''${pinch_args[@]}"

      # Build
      in_tmpdir hydrate ${config.system.build.nixos-rebuild}/bin/nixos-rebuild build
      while read user;do
        pushd /
        hydrate /run/wrappers/bin/sudo -u "$user" \
          ${pkgs.nix}/bin/nix-build --no-out-link '<nixpkgs>' -A "$(userenv_query "$user" .package)"
        popd
      done < <( config_query '.userEnvironments | keys []' )
      sync

      # Install
      hydrate ${config.system.build.nixos-rebuild}/bin/nixos-rebuild ${cfg.operation}
      sync
      while read user;do
        configured_action=$(userenv_query "$user" .otherPackagesAction)
        action=
        case "$configured_action" in
          keep)       action=-i;;
          remove)     action=-ir;;
          set-remove) action=--set;;
          *)          die "Unexpected otherPackagesAction $configured_action for user $user"
        esac
        hydrate /run/wrappers/bin/sudo -u "$user" \
          ${pkgs.nix}/bin/nix-env -f '<nixpkgs>' "$action" -A "$(userenv_query "$user" .package)"
        sync
      done < <( config_query '.userEnvironments | keys []' )
    ''}
  '';

in
{
  options = {
    system.autoUpgradeWithPinch = {

      enable = mkOption {
        type = types.bool;
        default = false;
        description = ''
          Whether to periodically upgrade NixOS to the latest version.
          Presumes that /etc/nixos is a git repo with a remote and
          contains a pinch file called "channels".
        '';
      };

      dates = mkOption {
        default = "04:40";
        type = types.str;
        description = ''
          Specification (in the format described by
          <citerefentry><refentrytitle>systemd.time</refentrytitle>
          <manvolnum>7</manvolnum></citerefentry>) of the time at
          which the update will occur.
        '';
      };

      operation = lib.mkOption {
        type = lib.types.enum [
          "switch"
          "boot"
        ];
        default = "switch";
        example = "boot";
        description = ''
          Whether to run
          `nixos-rebuild switch --upgrade` or run
          `nixos-rebuild boot --upgrade`
        '';
      };

      signingKeys = mkOption {
        type = types.listOf types.path;
        description = ''
          Files containing GPG keys that are authorized to sign updates.
          Updates are only merged if the commit at the tip of the remote
          ref is signed with one of these keys.
        '';
      };

      upgradeConfig = mkOption {
        type = types.listOf types.path;
        description = ''
          Configuration files that specify what git repo paths to pull, what
          pinch files to update from, and what user environments to update.
          These are specified in separate configuration files processed at
          update time so that changes to this configuration take effect in
          the same update cycle.
        '';
      };

      upgradeConfigOwnershipPolicy = mkOption {
        type = types.enum [
          "root"
          "wheel"
          "any"
        ];
        default = "root";
        description = ''
          Verify ownership of upgrade config files before using them for
          system upgrades.

          root = Config must be writable only by root.
          wheel = Config must be writable only by root and wheel.
          any = No checks.  Not recommended.
        '';
      };
    };
  };

  config = lib.mkIf cfg.enable {

    security.sudo.extraRules = lib.mkAfter [
      {
        groups = [ "users" ];
        runAs = "root";
        commands = [
          {
            command = "${auto-upgrade-script}";
            options = [
              "NOPASSWD"
              "NOSETENV"
            ];
          }
        ];
      }
    ];
    # NOSETENV above still allows through ~17 vars, including PATH.  Block those
    # as well:
    security.sudo.extraConfig = ''
      Defaults!${auto-upgrade-script} !env_check
      Defaults!${auto-upgrade-script} !env_keep
    '';

    nixpkgs.overlays = [
      (import ../overlays/pinch.nix)
      (import ../overlays/polite-merge.nix)
      (self: super: {
        auto-upgrade = super.writeShellScriptBin "auto-upgrade" ''
          /run/wrappers/bin/sudo ${auto-upgrade-script}
        '';
      })
    ];

    environment.systemPackages = [ pkgs.auto-upgrade ];

    systemd.services.nixos-upgrade = {
      description = "NixOS Upgrade";
      restartIfChanged = false;
      unitConfig.X-StopOnRemoval = false;
      serviceConfig.Type = "oneshot";
      environment =
        config.nix.envVars
        // {
          inherit (config.environment.sessionVariables) NIX_PATH;
          HOME = "/root";
        }
        // config.networking.proxy.envVars;

      path = with pkgs; [
        config.nix.package.out
        coreutils
        git
        gitMinimal
        gnutar
        gzip
        xz.bin
      ];

      script = ''
        set -eo pipefail

        # Chill for awhile before applying updates.  If applying an update
        # badly breaks things, we want a window in which an operator can
        # intervene either to fix the problem or disable automatic updates.
        sleep 2h

        ${auto-upgrade-script}
      '';

      startAt = cfg.dates;
    };
  };
}
