From 6d2af3720335c3f3ca4a56672150e48aa20c58f8 Mon Sep 17 00:00:00 2001
From: sp3nx0r <sp3nx0r@gmail.com>
Date: Wed, 2 Oct 2019 14:05:22 -0500
Subject: [PATCH] fixing getopts, adding host specific files to backup,
 refactoring

---
 scripts/update_pwnagotchi.sh | 58 +++++++++++++++++++++---------------
 1 file changed, 34 insertions(+), 24 deletions(-)

diff --git a/scripts/update_pwnagotchi.sh b/scripts/update_pwnagotchi.sh
index 0039e7bc..f68d9253 100644
--- a/scripts/update_pwnagotchi.sh
+++ b/scripts/update_pwnagotchi.sh
@@ -13,13 +13,14 @@ function usage() {
     cat <<EOF
 
  usage: $0 [OPTIONS]
- 
+ Note: This should be run from the pwnagotchi itself!
+
    Options:
       -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.)
+      -b        # Backup the current pwnagotchi config, then overwrite with defaults.
+      -r        # Restore the current pwnagotchi config after upgrade. (-b will be enabled.)
       -h        # Shows this help.
 
 EOF
@@ -36,16 +37,11 @@ function test_root() {
 function test_github() {
     wget -q  --spider $GIT_URL
     if [ $? -ne 0 ]; then
-        echo "[!] Cannot reach github. This script requires internet access, ensure connection sharing is working."
+        echo "[!] Cannot reach github. This script requires internet access and DNS resolution, ensure connection sharing is working and valid DNS server in /etc/resolv.conf."
         exit 2
     fi
 }
-
-echo "[+] Checking prerequisites."
-test_root
-test_github
-
-while getopts ":v:u:m:b:r:h" o; do
+while getopts "v:u:m:brh" o; do
   case "${o}" in
     v)
       VERSION="${OPTARG}"
@@ -53,22 +49,22 @@ while getopts ":v:u:m:b:r:h" o; do
     u)
       GIT_URL="${OPTARG}"
       ;;
+    b)
+      BACKUPCONFIG=1
+      ;;
+    r)
+      BACKUPCONFIG=1
+      RESTORECONFIG=1
+      ;;
+    h)
+      usage
+      ;;
     m)
       if [[ "${SUPPORTED_RESTART_MODES[*]}" =~ ${OPTARG} ]]; then
         MODE="${OPTARG}"
       else
         usage
-      fi      
-      ;;      
-    b)
-      BACKUPCONFIG=1
-      ;;
-    r)
-      BACKUPCONFIG=1    
-      RESTORECONFIG=1
-      ;;      
-    h)
-      usage
+      fi
       ;;
     *)
       usage
@@ -77,29 +73,43 @@ while getopts ":v:u:m:b:r:h" o; do
 done
 shift $((OPTIND-1))
 
+echo "[+] Checking prerequisites."
+test_root
+test_github
+
 # clean up old files, clone master, set checkout to commit if needed.
 echo "[+] Cloning to $GIT_FOLDER..."
-rm $GIT_FOLDER -rf
+#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
+    echo "[+] Creating backup of host files"
+    mv /etc/hosts /etc/hosts.bak -f
+    mv /etc/hostname /etc/hostname.bak -f
+    mv /etc/motd /etc/motd.bak -f
 fi
+
+echo "[+] Installing $(git log -1 --format="%h")"
 rm /root/pwnagotchi -rf # ensures old files are removed
 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
+    echo "[+] Restoring backup of host files"
+    mv /etc/hosts.bak /etc/hosts -f
+    mv /etc/hostname.bak /etc/hostname -f
+    mv /etc/motd.bak /etc/motd -f
 fi
 
 echo "[+] Restarting pwnagotchi in $MODE mode. $( screen -X -S pwnagotchi quit)"
@@ -108,4 +118,4 @@ if [ $MODE == "auto" ]; then
 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
+echo "[+] Finished"