From bc2bca79b5139d4d65d3fd7e824ff2793e12cdda Mon Sep 17 00:00:00 2001
From: Justin-p <justin-p@users.noreply.github.com>
Date: Wed, 2 Oct 2019 10:40:16 +0200
Subject: [PATCH] updated arguments to getopts. Added -m for restart mode.
 Moved prerequisites checks.

---
 scripts/update_pwnagotchi.sh | 119 ++++++++++++++++++++++-------------
 1 file changed, 74 insertions(+), 45 deletions(-)

diff --git a/scripts/update_pwnagotchi.sh b/scripts/update_pwnagotchi.sh
index 60107214..34e4eedf 100644
--- a/scripts/update_pwnagotchi.sh
+++ b/scripts/update_pwnagotchi.sh
@@ -1,82 +1,111 @@
 #!/bin/bash
 # Default variables
-git_folder="/tmp/pwnagotchi"
-git_url="https://github.com/evilsocket/pwnagotchi/"
-version="master"
-backupconfig=0
-restoreconfig=0
+GIT_FOLDER="/tmp/pwnagotchi"
+GIT_URL="https://github.com/evilsocket/pwnagotchi/"
+VERSION="master"
+SUPPORTED_RESTART_MODES=( 'auto' 'manual' )
+MODE="auto"
+BACKUPCONFIG=0
+RESTORECONFIG=0
 
 # Functions
 function usage() {
-        cat <<EOF
+    cat <<EOF
 
  usage: $0 [OPTIONS]
-
+ 
    Options:
-      -v, --version              # Version to update to, can be a branch or commit. (default: master)
-      -u, --url                  # Url to clone from. (default: https://github.com/evilsocket/pwnagotchi)
-      -bc, --backupconfig        # Backup the current pwnagotchi config.
-      -rc, --restoreconfig       # Restore the current pwnagotchi config. -bc will be enabled.
-      -h, --help                 # Shows this help.
+      -v        # Version to update to, can be a branch or commit. (default: master)
+      -u        # Url to clone from. (default: https://github.com/evilsocket/pwnagotchi)
+      -m        # Mode to restart to. (Supported: ${SUPPORTED_RESTART_MODES[*]}; default: auto)
+      -b        # Backup the current pwnagotchi config.
+      -r        # Restore the current pwnagotchi config. (-b will be enabled.)
+      -h        # Shows this help.
 
 EOF
-        exit 0
+    exit 0
 }
 
 function test_root() {
     if ! [ $(id -u) = 0 ]; then
-            echo "[!] This script must be run as root."
-            exit 1
+        echo "[!] This script must be run as root."
+        exit 1
     fi
 }
 
 function test_github() {
-    wget -q  --spider $git_url
+    wget -q  --spider $GIT_URL
     if [ $? -ne 0 ]; then
-            echo "[!] Cannot reach github. This script requires internet access, ensure connection sharing is working."
-            exit 2
+        echo "[!] Cannot reach github. This script requires internet access, ensure connection sharing is working."
+        exit 2
     fi
 }
 
-# Commandline arguments
-while [[ "$#" -gt 0 ]]; do case $1 in
-        -v|--version) version="$2"; shift;;
-        -u|--url) git_url="$2"; shift;;
-        -bc|--backupconfig) backupconfig=1; shift;;
-        -rc|--restoreconfig) backupconfig=1 restoreconfig=1; shift;;
-        -h|--help) usage;;
-        *) echo "Unknown parameter passed: $1"; exit 3;;
-esac; shift; done
-
 echo "[+] Checking prerequisites."
 test_root
 test_github
 
+while getopts ":v:u:m:b:r:h" o; do
+  case "${o}" in
+    v)
+      VERSION="${OPTARG}"
+      ;;
+    u)
+      GIT_URL="${OPTARG}"
+      ;;
+    m)
+      if [[ "${SUPPORTED_RESTART_MODES[*]}" =~ ${OPTARG} ]]; then
+        MODE="${OPTARG}"
+      else
+        usage
+      fi      
+      ;;      
+    b)
+      BACKUPCONFIG=1
+      ;;
+    r)
+      BACKUPCONFIG=1    
+      RESTORECONFIG=1
+      ;;      
+    h)
+      usage
+      ;;
+    *)
+      usage
+      ;;
+  esac
+done
+shift $((OPTIND-1))
+
 # clean up old files, clone master, set checkout to commit if needed.
-echo "[+] Cloning to $git_folder..."
-rm $git_folder -rf
-git clone $git_url $git_folder -q
-cd $git_folder
-if [ $version != "master" ]; then
-        git checkout $version -q
+echo "[+] Cloning to $GIT_FOLDER..."
+rm $GIT_FOLDER -rf
+git clone $GIT_URL $GIT_FOLDER -q
+cd $GIT_FOLDER
+if [ $VERSION != "master" ]; then
+    git checkout $VERSION -q
 fi
 echo "[+] Installing $(git log -1 --format="%h")"
 
 echo "[+] Updating..."
-if [ $backupconfig -eq 1 ]; then
-        echo "[+] Creating backup of config.yml"
-        mv /root/pwnagotchi/config.yml /root/config.yml.bak -f
+if [ $BACKUPCONFIG -eq 1 ]; then
+    echo "[+] Creating backup of config.yml"
+    mv /root/pwnagotchi/config.yml /root/config.yml.bak -f
 fi
 rm /root/pwnagotchi -rf # ensures old files are removed
-rsync -aPq $git_folder/sdcard/boot/*   /boot/
-rsync -aPq $git_folder/sdcard/rootfs/* /
+rsync -aPq $GIT_FOLDER/sdcard/boot/*   /boot/
+rsync -aPq $GIT_FOLDER/sdcard/rootfs/* /
 cd /tmp
-rm $git_folder -rf
-if [ $restoreconfig -eq 1 ]; then
-        echo "[+] Restoring backup of config.yml"
-        mv /root/config.yml.bak /root/pwnagotchi/config.yml -f
+rm $GIT_FOLDER -rf
+if [ $RESTORECONFIG -eq 1 ]; then
+    echo "[+] Restoring backup of config.yml"
+    mv /root/config.yml.bak /root/pwnagotchi/config.yml -f
 fi
 
-echo "[+] Restarting pwnagotchi in auto mode. $( screen -X -S pwnagotchi quit)"
-sudo -H -u root /usr/bin/screen -dmS pwnagotchi -c /root/pwnagotchi/data/screenrc.auto
+echo "[+] Restarting pwnagotchi in $MODE mode. $( screen -X -S pwnagotchi quit)"
+if [ $MODE == "auto" ]; then
+    sudo -H -u root /usr/bin/screen -dmS pwnagotchi -c /root/pwnagotchi/data/screenrc.auto
+elif [ $MODE == "manual" ]; then
+    sudo -H -u root /usr/bin/screen -dmS pwnagotchi -c /root/pwnagotchi/data/screenrc.manual
+fi
 echo "[+] Finished"
\ No newline at end of file