# 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 []' )

        # Install
        hydrate ${config.system.build.nixos-rebuild}/bin/nixos-rebuild switch
        while read user;do
          remove_arg=-r
          if [[ "$(userenv_query "$user" .otherPackagesAction)" == keep ]];then
            remove_arg=
          fi
          hydrate /run/wrappers/bin/sudo -u "$user" \
            ${pkgs.nix}/bin/nix-env -f '<nixpkgs>' $remove_arg -iA "$(userenv_query "$user" .package)"
        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.
        '';
      };

      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" ];
      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;
    };
  };
}