diff --git a/bin/auto-update-nixos b/bin/auto-update-nixos
index 05f812b..7d36b4a 100755
--- a/bin/auto-update-nixos
+++ b/bin/auto-update-nixos
@@ -7,6 +7,7 @@ hostname=$(/run/current-system/sw/bin/hostname) # The name of the host to build
 flakeDir="${FLAKE_DIR}"                         # Path to the flake file (and optionally the hostname). Defaults to the FLAKE_DIR environment variable.
 update=false                                    # Whether to update flake.lock (false by default)
 user=$(/run/current-system/sw/bin/whoami)       # Which user account to use for git commands (defaults to whoever called the script)
+reboot=false
 remainingArgs=""                                # All remaining arguments that haven't yet been processed (will be passed to nixos-rebuild)
 
 function usage() {
@@ -42,6 +43,10 @@ while [[ $# -gt 0 ]]; do
 		update=true
 		shift
 		;;
+	--reboot|-r)
+		reboot=true
+		shift
+		;;
 	--operation|-o)
 		operation="$2"
 		shift
@@ -72,7 +77,11 @@ fi
 
 cd $flakeDir
 
+current_branch=$(git branch --show-current)
+[ "$current_branch" != "master" ] && echo "Not on master branch. Aborting auto-update" && exit 0
+
 echo "Pulling the latest version of the repository..."
+/run/wrappers/bin/sudo -u $user git stash
 /run/wrappers/bin/sudo -u $user git pull
 
 if [ $update = true ]; then
@@ -88,11 +97,14 @@ echo "Running this operation: nixos-rebuild $operation $options"
 /run/wrappers/bin/sudo -u root /run/current-system/sw/bin/nixos-rebuild $operation $options
 
 echo "Checking if reboot is necessary"
-reboot=$(diff <(readlink /run/booted-system/{initrd,kernel,kernel-modules}) <(readlink /nix/var/nix/profiles/system/{initrd,kernel,kernel-modules}))
-if [ -z "$reboot" ]; then
-  echo "No reboot necessary, exiting."
-  exit 0
-else
+reboot_diff=$(diff <(readlink /run/booted-system/{initrd,kernel,kernel-modules}) <(readlink /nix/var/nix/profiles/system/{initrd,kernel,kernel-modules}))
+if [ -n "$reboot_diff" ] && [ $reboot == true ]; then
   echo "System requires a reboot. Rebooting now..."
   reboot
+else
+  echo "No reboot necessary."
+  echo "Update complete."
+  exit 0
 fi
+echo "Update complete."
+exit 0