From 34bafe51ffe6a47837c8bcde104fb89e4038f764 Mon Sep 17 00:00:00 2001
From: Dev Dua <devdua@live.com>
Date: Thu, 12 Jul 2018 13:39:10 +0530
Subject: [PATCH] Added license option & improved update process

---
 helk_update.sh | 79 +++++++++++++++++++++++++++++++++++---------------
 1 file changed, 56 insertions(+), 23 deletions(-)

diff --git a/helk_update.sh b/helk_update.sh
index afd1dd2..bdb8f29 100755
--- a/helk_update.sh
+++ b/helk_update.sh
@@ -2,16 +2,35 @@
 
 # HELK script: helk_update.sh
 # HELK script description: Update and Rebuild HELK
-# HELK build version: v0.1.1-alpha07062018
+# HELK build Stage: Alpha
 # HELK ELK version: 6.3.1
+# Author: Roberto Rodriguez (@Cyb3rWard0g)
 # Script Author: Dev Dua (@devdua)
-# License: BSD 3-Clause
+# License: GPL-3.0
 
 if [[ $EUID -ne 0 ]]; then
    echo "[HELK-UPDATE-INFO] YOU MUST BE ROOT TO RUN THIS SCRIPT!!!" 
    exit 1
 fi
 
+set_helk_license(){    
+    # *********** Accepting Defaults or Allowing user to set HELK License ***************
+    local license_input
+    read -t 30 -p "[HELK-UPDATE-INFO] Set HELK License. Default value is basic: " -e -i "basic" license_input
+    license_choice=${license_input:-"basic"}
+    # *********** Validating License Input ***************
+    case $license_choice in
+        basic)
+        ;;
+        trial)
+        ;;
+        *)
+            echo "[HELK-UPDATE-ERROR] Not a valid license. Valid Options: basic or trial"
+            exit 1
+        ;;
+    esac
+}
+
 check_min_requirements(){
     systemKernel="$(uname -s)"
     echo "[HELK-UPDATE-INFO] HELK being hosted on a $systemKernel box"
@@ -47,31 +66,31 @@ check_github(){
         echo "Git successfully installed." >> $LOGFILE
     fi
 
-    if HELK_REPO_EXISTS=$(git remote | grep "helk-repo") && [ -z "$HELK_REPO_EXISTS" ]; then
+    if [[ -z "$(git remote | grep helk-repo)" ]]; then
         git remote add helk-repo https://github.com/Cyb3rWard0g/HELK.git  >> $LOGFILE 2>&1 
     else
         echo "HELK repo exists" >> $LOGFILE 2>&1
     fi
 
-    git checkout master >> $LOGFILE 2>&1
-
-    if PULL_NEEDED=$(git rev-list --left-right --count origin/master...helk-repo/master | awk '{print $2}') && [ ! "$PULL_NEEDED" == "0" ]; then
-        echo "New release available. Commit diff --> $PULL_NEEDED" >> $LOGFILE 2>&1
-        echo "[HELK-UPDATE-INFO] New release available. Pulling new code."
-        git pull helk-repo master >> $LOGFILE 2>&1
-        ERROR=$?
-        if [ $ERROR -ne 0 ]; then
-            echo "[!] Could not run pull latest code (Error Code: $ERROR). Check $LOGFILE for details."
-            exit 1
+    if COMMIT_DIFF=$(git rev-list --count master...helk-repo/master) && [ ! "$COMMIT_DIFF" == "0" ]; then
+        echo "Possibly new release available. Commit diff --> $COMMIT_DIFF" >> $LOGFILE 2>&1
+        IS_MASTER_BEHIND=$(git branch -v | grep master | grep behind)
+        
+        if [[ -z $IS_MASTER_BEHIND ]]; then
+            echo "Current master branch ahead of remote branch. Exiting..." >> $LOGFILE 2>&1
+            echo "[HELK-UPDATE-INFO] No updates available."
+        else            
+            echo "[HELK-UPDATE-INFO] New release available. Pulling new code."
+            git checkout master >> $LOGFILE 2>&1
+            git pull helk-repo master >> $LOGFILE 2>&1
+            REBUILD_NEEDED=1
         fi
-        REBUILD_NEEDED=1
     else
         echo "[HELK-UPDATE-INFO] No updates available."    
     fi
 }
 
 get_jupyter_token(){
-    echo "[HELK-UPDATE-INFO] Waiting for HELK services and Jupyter Server to start.."
     until curl -s localhost:8880 -o /dev/null; do
         sleep 1
     done
@@ -82,24 +101,38 @@ get_jupyter_token(){
 LOGFILE="/var/log/helk-update.log"
 REBUILD_NEEDED=0
 
+echo "[HELK-UPDATE-INFO] Checking GitHub for updates..."   
 check_github
 
 if [ $REBUILD_NEEDED == 1 ]; then
+    set_helk_license
     echo -e "[HELK-UPDATE-INFO] Stopping HELK and starting update"
-    docker-compose down >> $LOGFILE 2>&1
-    check_min_requirements
-
-    echo "[HELK-UPDATE-INFO] Rebuilding HELK via docker-compose"
-    docker-compose up --build -d -V --force-recreate --always-recreate-deps >> $LOGFILE 2>&1
+    docker-compose -f docker-compose-elk-${license_choice}.yml down >> $LOGFILE 2>&1
     ERROR=$?
     if [ $ERROR -ne 0 ]; then
-        echo "[!] Could not run HELK via docker-compose (Error Code: $ERROR). Check $LOGFILE for details."
+        echo -e "[!] Could not stop HELK via docker-compose (Error Code: $ERROR). You're possibly running a different HELK license than chosen - $license_choice"
         exit 1
     fi
 
-    sleep 180
+    check_min_requirements
+
+    echo "[HELK-UPDATE-INFO] Rebuilding HELK via docker-compose"
+    docker-compose -f docker-compose-elk-${license_choice}.yml up --build -d -V --force-recreate --always-recreate-deps >> $LOGFILE 2>&1
+    ERROR=$?
+    if [ $ERROR -ne 0 ]; then
+        echo -e "[!] Could not run HELK via docker-compose (Error Code: $ERROR). Check $LOGFILE for details."
+        exit 1
+    fi
+    
+    secs=$((3 * 60))
+    while [ $secs -gt 0 ]; do
+        echo -ne "\033[0K\r[HELK-UPDATE-INFO] Rebuild succeeded, waiting $secs seconds for services to start"
+        sleep 1
+        : $((secs--))
+    done
+
     get_jupyter_token
     echo -e "[HELK-UPDATE-INFO] YOUR HELK HAS BEEN UPDATED!"
 else
     echo -e "[HELK-UPDATE-INFO] YOUR HELK IS ALREADY UP-TO-DATE."
-fi
+fi
\ No newline at end of file