From 6856d12a22d0d942f8153dc2b9e17e8897d3a912 Mon Sep 17 00:00:00 2001
From: UGA Innovation Factory <engr-ugaif@uga.edu>
Date: Fri, 19 Dec 2025 10:40:34 -0500
Subject: [PATCH] update the updater to support flags for privileged users

---
 sw/updater.nix | 148 +++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 136 insertions(+), 12 deletions(-)

diff --git a/sw/updater.nix b/sw/updater.nix
index 3d169e1..445d766 100644
--- a/sw/updater.nix
+++ b/sw/updater.nix
@@ -35,25 +35,149 @@ with lib;
       (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"
 
-        # Start following logs in the background
-        journalctl -fu "$UNIT" -n 0 --output=cat &
-        JPID=$!
+        # 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
 
-        # Start the service and wait for it to finish
-        if systemctl start --wait --no-ask-password "$UNIT"; then
-          STATUS=$?
-        else
-          STATUS=$?
+        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
 
-        sleep 2
+        # 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
 
-        # Kill the log follower
-        kill "$JPID" 2>/dev/null || true
+        # Parse flags
+        while [ "$#" -gt 0 ]; do
+          case "$1" in
+            --local-repo)
+              REPO_MODE="local"
+              LOCAL_PATH="''${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
 
-        exit "$STATUS"
+        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="''${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
+
+        exec nixos-rebuild switch --refresh --print-build-logs $impureFlag --flake "''${flakeRef}"
       '')
     ];
 
-- 
2.39.5