athenix/sw/updater.nix

{
  config,
  lib,
  pkgs,
  ...
}:

with lib;

{
  options.athenix.sw.remoteBuild = lib.mkOption {
    type = types.submodule {
      options = {
        hosts = mkOption {
          type = types.listOf types.str;
          default = [ "engr-ugaif@192.168.11.133 x86_64-linux" ];
          description = "List of remote build hosts for system rebuilding.";
        };

        enable = mkOption {
          type = types.bool;
          default = false;
          description = "Whether to enable remote build for 'update-system' command.";
        };
      };
    };
    default = { };
    description = "Remote build configuration";
  };

  config = {
    athenix.sw.remoteBuild.enable = lib.mkDefault (config.athenix.sw.type == "tablet-kiosk");

    environment.systemPackages = [
      (pkgs.writeShellScriptBin "update-system" ''
        set -euo pipefail

        RED='\033[31m'; NC='\033[0m'

        is_root() { [ "''${EUID:-$(id -u)}" -eq 0 ]; }
        in_wheel() { id -nG 2>/dev/null | tr ' ' '\n' | grep -qx wheel; }

        # Service path for unprivileged (no flags)
        UNIT="update-system.service"

        # Figure out the "real" invoking user, even under sudo.
        INVOKER_USER="''${SUDO_USER:-$(id -un)}"
        INVOKER_HOME="$(getent passwd "$INVOKER_USER" | cut -d: -f6)"
        if [ -z "$INVOKER_HOME" ]; then
          # fallback if getent is weird in some containers
          INVOKER_HOME="''${HOME:-/home/$INVOKER_USER}"
        fi

        # Defaults for flagged mode
        DEFAULT_REMOTE_URL="https://git.factory.uga.edu/UGA-Innovation-Factory/athenix"
        REPO_MODE="default"     # default | local | remote
        LOCAL_PATH=""
        REMOTE_URL=""
        BRANCH=""
        IMPURE=0

        usage() {
          cat >&2 <<'EOF'
        usage:
          update-system
          update-system [--local-repo[=PATH]] [--remote-repo=URL] [--branch=BRANCH] [--impure]

        notes:
          - No flags: runs the systemd service (works for unprivileged users via polkit).
          - Any flags: only allowed for root or wheel (runs nixos-rebuild directly).
        EOF
          exit 2
        }

        # No flags -> polkit-friendly systemd service route
        if [ "$#" -eq 0 ]; then
          journalctl -fu "$UNIT" -n 0 --output=cat &
          JPID=$!

          if systemctl start --wait --no-ask-password "$UNIT"; then
            STATUS=$?
          else
            STATUS=$?
          fi

          sleep 2
          kill "$JPID" 2>/dev/null || true
          exit "$STATUS"
        fi

        # Flags -> require root or wheel
        if ! is_root && ! in_wheel; then
          printf "''${RED}error:''${NC} flags are only allowed for root or wheel. Run without flags (service path), or use sudo / add yourself to wheel.\n" >&2
          exit 2
        fi

        # Parse flags
        while [ "$#" -gt 0 ]; do
          case "$1" in
            --local-repo)
              REPO_MODE="local"
              LOCAL_PATH="$INVOKER_HOME/athenix"
              shift
              ;;
            --local-repo=*)
              REPO_MODE="local"
              LOCAL_PATH="''${1#*=}"
              shift
              ;;
            --remote-repo=*)
              REPO_MODE="remote"
              REMOTE_URL="''${1#*=}"
              shift
              ;;
            --branch)
              [ "$#" -ge 2 ] || usage
              BRANCH="$2"
              shift 2
              ;;
            --branch=*)
              BRANCH="''${1#*=}"
              shift
              ;;
            --impure)
              IMPURE=1
              shift
              ;;
            -h|--help) usage ;;
            *)
              printf "''${RED}error:''${NC} unknown argument: %s\n" "$1" >&2
              usage
              ;;
          esac
        done

        if [ "$REPO_MODE" = "local" ] && [ -n "$REMOTE_URL" ]; then
          printf "''${RED}error:''${NC} can't use --local-repo and --remote-repo together.\n" >&2
          exit 2
        fi

        host="''${HOSTNAME:-$(hostname)}"

        # Build flake ref
        if [ "$REPO_MODE" = "local" ]; then
          [ -n "$LOCAL_PATH" ] || LOCAL_PATH="$INVOKER_HOME/athenix"

          # Clone default repo if missing
          if [ ! -d "$LOCAL_PATH" ]; then
            printf "local repo not found at %s, cloning %s...\n" "$LOCAL_PATH" "$DEFAULT_REMOTE_URL" >&2
            if [ -n "$BRANCH" ]; then
              git clone --branch "$BRANCH" "$DEFAULT_REMOTE_URL" "$LOCAL_PATH"
            else
              git clone "$DEFAULT_REMOTE_URL" "$LOCAL_PATH"
            fi
          fi

          flakeRef="''${LOCAL_PATH}#''${host}"
        else
          url="''${REMOTE_URL:-$DEFAULT_REMOTE_URL}"

          if echo "$url" | grep -qE '^(https?|ssh)://'; then
            base="git+''${url}"
          elif echo "$url" | grep -qE '^[^/@:]+@[^/:]+:'; then
            # scp-style: git@host:owner/repo(.git)
            userhost="''${url%%:*}"
            path="''${url#*:}"
            base="git+ssh://''${userhost}/''${path}"
          else
            base="''${url}"
          fi

          if [ -n "$BRANCH" ]; then
            if echo "$base" | grep -q '?'; then
              base="''${base}&ref=''${BRANCH}"
            else
              base="''${base}?ref=''${BRANCH}"
            fi
          fi

          flakeRef="''${base}#''${host}"
        fi

        impureFlag=""
        if [ "$IMPURE" -eq 1 ]; then
          impureFlag="--impure"
        fi

        # If not root, re-exec via sudo to do the actual switch.
        # Preserve our computed invoker context so sudo doesn't "helpfully" change it.
        if ! is_root; then
          exec sudo --preserve-env=HOME,USER,LOGNAME \
            nixos-rebuild switch --refresh --print-build-logs $impureFlag --flake "$flakeRef"
        else
          exec nixos-rebuild switch --refresh --print-build-logs $impureFlag --flake "$flakeRef"
        fi
      '')
    ];

    systemd.services.update-system = {
      enable = true;
      description = "System daemon to one-shot run the Nix updater from fleet flake as root";
      path = with pkgs; [
        git
        nixos-rebuild
        nix
      ];
      serviceConfig = {
        Type = "oneshot";
        ExecStart =
          let
            hosts = config.athenix.sw.remoteBuild.hosts;
            builders = lib.strings.concatMapStringsSep ";" (x: x) hosts;
            rebuildCmd = "${pkgs.nixos-rebuild}/bin/nixos-rebuild switch --refresh";
            source = "--flake git+https://git.factory.uga.edu/UGA-Innovation-Factory/athenix";
            remoteBuildFlags =
              if config.athenix.sw.remoteBuild.enable then ''--builders "${builders}"'' else "";
          in
          "${rebuildCmd} ${remoteBuildFlags} --print-build-logs ${source}#${config.networking.hostName}";
        User = "root";
        Group = "root";
        TimeoutStartSec = "0";
      };
    };

    security.polkit = {
      debug = true;
      enable = true;
      extraConfig = ''
        polkit.addRule(function(action, subject) {
          if (action.id == "org.freedesktop.systemd1.manage-units" &&
              action.lookup("unit") == "update-system.service" &&
              action.lookup("verb") == "start" &&
              subject.isInGroup("users")) {
            return polkit.Result.YES;
          }
        });
      '';
    };
  };
}