#!/bin/bash

# Copyright 2014,2015,2016,2017,2018,2019,2020 Security Onion Solutions, LLC

#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

. /usr/sbin/so-common
UPDATE_DIR=/tmp/sogh/securityonion
INSTALLEDVERSION=$(cat /etc/soversion)
INSTALLEDSALTVERSION=$(salt --versions-report | grep Salt: | awk {'print $2'})
DEFAULT_SALT_DIR=/opt/so/saltstack/default
BATCHSIZE=5
SOUP_LOG=/root/soup.log
exec 3>&1 1>${SOUP_LOG} 2>&1

manager_check() {
  # Check to see if this is a manager
  MANAGERCHECK=$(cat /etc/salt/grains | grep role | awk '{print $2}')
  if [[ "$MANAGERCHECK" =~ ^('so-eval'|'so-manager'|'so-standalone'|'so-managersearch'|'so-import')$ ]]; then
    echo "This is a manager. We can proceed."
    MINIONID=$(salt-call grains.get id --out=txt|awk -F: {'print $2'}|tr -d ' ')
  else
    echo "Please run soup on the manager. The manager controls all updates."
    exit 0
  fi
}

airgap_mounted() {
  # Let's see if the ISO is already mounted. 
  if [ -f /tmp/soagupdate/SecurityOnion/VERSION ]; then
    echo "The ISO is already mounted"
  else
    echo ""
    echo "Looks like we need access to the upgrade content"
    echo "" 
    echo "If you just copied the .iso file over you can specify the path."
    echo "If you burned the ISO to a disk the standard way you can specify the device."
    echo "Example: /home/user/securityonion-2.X.0.iso"
    echo "Example: /dev/sdx1"
    echo ""
    read -p 'Enter the location of the iso: ' ISOLOC
    if [ -f $ISOLOC ]; then
      # Mounting the ISO image
      mkdir -p /tmp/soagupdate
      mount -t iso9660 -o loop $ISOLOC /tmp/soagupdate
      # Make sure mounting was successful
      if [ ! -f /tmp/soagupdate/SecurityOnion/VERSION ]; then
        echo "Something went wrong trying to mount the ISO."
        echo "Ensure you verify the ISO that you downloaded."
        exit 0
      else
        echo "ISO has been mounted!"
      fi  
    elif [ -f $ISOLOC/SecurityOnion/VERSION ]; then
      ln -s $ISOLOC /tmp/soagupdate
      echo "Found the update content"
    else 
      mkdir -p /tmp/soagupdate
      mount $ISOLOC /tmp/soagupdate
      if [ ! -f /tmp/soagupdate/SecurityOnion/VERSION ]; then
        echo "Something went wrong trying to mount the device."
        echo "Ensure you verify the ISO that you downloaded."
        exit 0
      else
        echo "Device has been mounted!"
      fi        
    fi
  fi
}

check_airgap() {
  # See if this is an airgap install
  AIRGAP=$(cat /opt/so/saltstack/local/pillar/global.sls | grep airgap | awk '{print $2}')
  if [[ "$AIRGAP" == "True" ]]; then
      is_airgap=0
      UPDATE_DIR=/tmp/soagupdate/SecurityOnion
      AGDOCKER=/tmp/soagupdate/docker
      AGREPO=/tmp/soagupdate/Packages
  else 
      is_airgap=1
  fi
}

clean_dockers() {
  # Place Holder for cleaning up old docker images
  echo "Trying to clean up old dockers."
  docker system prune -a -f

}

clone_to_tmp() {
  # TODO Need to add a air gap option
  # Clean old files
  rm -rf /tmp/sogh
  # Make a temp location for the files
  mkdir -p /tmp/sogh
  cd /tmp/sogh
  SOUP_BRANCH=""
  if [ -n "$BRANCH" ]; then
    SOUP_BRANCH="-b $BRANCH"
  fi
  git clone $SOUP_BRANCH https://github.com/Security-Onion-Solutions/securityonion.git
  cd /tmp
  if [ ! -f $UPDATE_DIR/VERSION ]; then
    echo "Update was unable to pull from github. Please check your internet."
    exit 0
  fi
}

copy_new_files() {
  # Copy new files over to the salt dir
  cd $UPDATE_DIR
  rsync -a salt $DEFAULT_SALT_DIR/
  rsync -a pillar $DEFAULT_SALT_DIR/
  chown -R socore:socore $DEFAULT_SALT_DIR/
  chmod 755 $DEFAULT_SALT_DIR/pillar/firewall/addfirewall.sh
  cd /tmp
}

detect_os() {
	# Detect Base OS
	echo "Determining Base OS." >> "$SOUP_LOG" 2>&1
	if [ -f /etc/redhat-release ]; then
		OS="centos"
	elif [ -f /etc/os-release ]; then
		OS="ubuntu"
	fi
	echo "Found OS: $OS" >> "$SOUP_LOG" 2>&1
}

highstate() {
  # Run a highstate.
  salt-call state.highstate -l info queue=True
}

masterlock() {
  echo "Locking Salt Master"
  if [[ "$INSTALLEDVERSION" =~ rc.1 ]]; then
    TOPFILE=/opt/so/saltstack/default/salt/top.sls
    BACKUPTOPFILE=/opt/so/saltstack/default/salt/top.sls.backup
    mv -v $TOPFILE $BACKUPTOPFILE
    echo "base:" > $TOPFILE
    echo "  $MINIONID:" >> $TOPFILE
    echo "    - ca" >> $TOPFILE
    echo "    - ssl" >> $TOPFILE
    echo "    - elasticsearch" >> $TOPFILE
  fi
}

masterunlock() {
  echo "Unlocking Salt Master"
  if [[ "$INSTALLEDVERSION" =~ rc.1 ]]; then
    mv -v $BACKUPTOPFILE $TOPFILE
  fi
}

playbook() {
  echo "Applying playbook settings"
  if [[ "$INSTALLEDVERSION" =~ rc.1 ]]; then
    salt-call state.apply playbook.OLD_db_init
    rm -f /opt/so/rules/elastalert/playbook/*.yaml
    so-playbook-ruleupdate >> /root/soup_playbook_rule_update.log 2>&1 &
  fi
}

pillar_changes() {
    # This function is to add any new pillar items if needed.
    echo "Checking to see if pillar changes are needed."
    
    [[ "$INSTALLEDVERSION" =~ rc.1 ]] && rc1_to_rc2
    [[ "$INSTALLEDVERSION" =~ rc.2 ]] && rc2_to_rc3
    [[ "$INSTALLEDVERSION" =~ rc.3 ]] && rc3_to_2.3.0

}

rc1_to_rc2() {

  # Move the static file to global.sls
  echo "Migrating static.sls to global.sls"
  mv -v /opt/so/saltstack/local/pillar/static.sls /opt/so/saltstack/local/pillar/global.sls >> "$SOUP_LOG" 2>&1
  sed -i '1c\global:' /opt/so/saltstack/local/pillar/global.sls >> "$SOUP_LOG" 2>&1

  # Moving baseurl from minion sls file to inside global.sls
  local line=$(grep '^  url_base:' /opt/so/saltstack/local/pillar/minions/$MINIONID.sls)
  sed -i '/^  url_base:/d' /opt/so/saltstack/local/pillar/minions/$MINIONID.sls;
  sed -i "/^global:/a \\$line" /opt/so/saltstack/local/pillar/global.sls;

  # Adding play values to the global.sls
  local HIVEPLAYSECRET=$(tr -dc 'a-zA-Z0-9' < /dev/urandom  | fold -w 20 | head -n 1)
  local CORTEXPLAYSECRET=$(tr -dc 'a-zA-Z0-9' < /dev/urandom  | fold -w 20 | head -n 1)
  sed -i "/^global:/a \\  hiveplaysecret: $HIVEPLAYSECRET" /opt/so/saltstack/local/pillar/global.sls;
  sed -i "/^global:/a \\  cortexplaysecret: $CORTEXPLAYSECRET" /opt/so/saltstack/local/pillar/global.sls;

  # Move storage nodes to hostname for SSL
  # Get a list we can use:
  grep -A1 searchnode /opt/so/saltstack/local/pillar/data/nodestab.sls | grep -v '\-\-' | sed '$!N;s/\n/ /' | awk '{print $1,$3}' | awk '/_searchnode:/{gsub(/\_searchnode:/, "_searchnode"); print}' >/tmp/nodes.txt
  # Remove the nodes from cluster settings
  while read p; do
  local NAME=$(echo $p | awk '{print $1}')
  local IP=$(echo $p | awk '{print $2}')
  echo "Removing the old cross cluster config for $NAME"
  curl -XPUT -H 'Content-Type: application/json' http://localhost:9200/_cluster/settings -d '{"persistent":{"cluster":{"remote":{"'$NAME'":{"skip_unavailable":null,"seeds":null}}}}}'
  done </tmp/nodes.txt
  # Add the nodes back using hostname
  while read p; do
      local NAME=$(echo $p | awk '{print $1}')
      local EHOSTNAME=$(echo $p | awk -F"_" '{print $1}')
	    local IP=$(echo $p | awk '{print $2}')
      echo "Adding the new cross cluster config for $NAME"
      curl -XPUT http://localhost:9200/_cluster/settings -H'Content-Type: application/json' -d '{"persistent": {"search": {"remote": {"'$NAME'": {"skip_unavailable": "true", "seeds": ["'$EHOSTNAME':9300"]}}}}}'
  done </tmp/nodes.txt

  INSTALLEDVERSION=rc.2

}

rc2_to_rc3() {

  # move location of local.rules
  cp /opt/so/saltstack/default/salt/idstools/localrules/local.rules /opt/so/saltstack/local/salt/idstools/local.rules
  
  if [ -f /opt/so/saltstack/local/salt/idstools/localrules/local.rules ]; then
    cat /opt/so/saltstack/local/salt/idstools/localrules/local.rules >> /opt/so/saltstack/local/salt/idstools/local.rules
  fi
  rm -rf /opt/so/saltstack/local/salt/idstools/localrules
  rm -rf /opt/so/saltstack/default/salt/idstools/localrules

  # Rename mdengine to MDENGINE
  sed -i "s/  zeekversion/  mdengine/g" /opt/so/saltstack/local/pillar/global.sls
  # Enable Strelka Rules
  sed -i "/  rules:/c\  rules: 1" /opt/so/saltstack/local/pillar/global.sls

  INSTALLEDVERSION=rc.3

}

rc3_to_2.3.0() {
  # Fix Tab Complete
  if [ ! -f /etc/profile.d/securityonion.sh ]; then
    echo "complete -cf sudo" > /etc/profile.d/securityonion.sh
  fi

  {
      echo "redis_settings:"
      echo "  redis_maxmemory: 827"
      echo "playbook:"
      echo "  api_key: de6639318502476f2fa5aa06f43f51fb389a3d7f" 
  } >> /opt/so/saltstack/local/pillar/global.sls

  sed -i 's/playbook:/playbook_db:/' /opt/so/saltstack/local/pillar/secrets.sls
  {
    echo "playbook_admin: $(tr -dc 'a-zA-Z0-9' < /dev/urandom | fold -w 20 | head -n 1)"
    echo "playbook_automation: $(tr -dc 'a-zA-Z0-9' < /dev/urandom | fold -w 20 | head -n 1)"
  } >> /opt/so/saltstack/local/pillar/secrets.sls
}

space_check() {
  # Check to see if there is enough space
  CURRENTSPACE=$(df -BG / | grep -v Avail | awk '{print $4}' | sed 's/.$//')
  if [ "$CURRENTSPACE" -lt "10" ]; then
      echo "You are low on disk space. Upgrade will try and clean up space.";
      clean_dockers
  else
      echo "Plenty of space for upgrading"
  fi
  
}

unmount_update() {
  cd /tmp
  umount /tmp/soagupdate
}

update_centos_repo() {
  # Update the files in the repo
  echo "Syncing new updates to /nsm/repo"
  rsync -av $AGREPO/* /nsm/repo/
  echo "Creating repo"
  createrepo /nsm/repo
}

update_dockers() {
  if [ $is_airgap -eq 0 ]; then
    # Let's copy the tarball
    if [ ! -f $AGDOCKER/registry.tar ]; then
      echo "Unable to locate registry. Exiting"
      exit 0
    else
      echo "Stopping the registry docker"
      docker stop so-dockerregistry
      docker rm so-dockerregistry
      echo "Copying the new dockers over"
      tar xvf $AGDOCKER/registry.tar -C /nsm/docker-registry/docker
    fi
  else
    # List all the containers
    if [ $MANAGERCHECK == 'so-import' ]; then
      TRUSTED_CONTAINERS=( \
      "so-idstools" \
      "so-nginx" \
      "so-filebeat" \
      "so-suricata" \
      "so-soc" \
      "so-elasticsearch" \
      "so-kibana" \
      "so-kratos" \
      "so-suricata" \
      "so-registry" \
      "so-pcaptools" \
      "so-zeek" )
    elif [ $MANAGERCHECK != 'so-helix' ]; then
      TRUSTED_CONTAINERS=( \
      "so-acng" \
      "so-thehive-cortex" \
      "so-curator" \
      "so-domainstats" \
      "so-elastalert" \
      "so-elasticsearch" \
      "so-filebeat" \
      "so-fleet" \
      "so-fleet-launcher" \
      "so-freqserver" \
      "so-grafana" \
      "so-idstools" \
      "so-influxdb" \
      "so-kibana" \
      "so-kratos" \
      "so-logstash" \
      "so-minio" \
      "so-mysql" \
      "so-nginx" \
      "so-pcaptools" \
      "so-playbook" \
      "so-redis" \
      "so-soc" \
      "so-soctopus" \
      "so-steno" \
      "so-strelka-frontend" \
      "so-strelka-manager" \
      "so-strelka-backend" \
      "so-strelka-filestream" \
      "so-suricata" \
      "so-telegraf" \
      "so-thehive" \
      "so-thehive-es" \
      "so-wazuh" \
      "so-zeek" )
    else
      TRUSTED_CONTAINERS=( \
      "so-filebeat" \
      "so-idstools" \
      "so-logstash" \
      "so-nginx" \
      "so-redis" \
      "so-steno" \
      "so-suricata" \
      "so-telegraf" \
      "so-zeek" )
    fi
    
# Download the containers from the interwebs
    for i in "${TRUSTED_CONTAINERS[@]}"
    do
      # Pull down the trusted docker image
      echo "Downloading $i:$NEWVERSION"
      docker pull --disable-content-trust=false docker.io/$IMAGEREPO/$i:$NEWVERSION
      # Tag it with the new registry destination
      docker tag $IMAGEREPO/$i:$NEWVERSION $HOSTNAME:5000/$IMAGEREPO/$i:$NEWVERSION
      docker push $HOSTNAME:5000/$IMAGEREPO/$i:$NEWVERSION 
    done
  fi
    echo "Add Registry back if airgap"
    if [ $is_airgap -eq 0 ]; then
      docker load -i $AGDOCKER/registry_image.tar
    fi 

}

update_version() {
  # Update the version to the latest
  echo "Updating the Security Onion version file."
  echo $NEWVERSION > /etc/soversion
  sed -i "/  soversion:/c\  soversion: $NEWVERSION" /opt/so/saltstack/local/pillar/global.sls
}

upgrade_check() {
    # Let's make sure we actually need to update.
    NEWVERSION=$(cat $UPDATE_DIR/VERSION)
    if [ "$INSTALLEDVERSION" == "$NEWVERSION" ]; then
      echo "You are already running the latest version of Security Onion."
      exit 0
    fi 
}

upgrade_check_salt() {
    NEWSALTVERSION=$(grep version: $UPDATE_DIR/salt/salt/master.defaults.yaml | awk {'print $2'})
    if [ "$INSTALLEDSALTVERSION" == "$NEWSALTVERSION" ]; then
      echo "You are already running the correct version of Salt for Security Onion."
    else
      UPGRADESALT=1
    fi
}   
upgrade_salt() {
      SALTUPGRADED=True
      echo "Performing upgrade of Salt from $INSTALLEDSALTVERSION to $NEWSALTVERSION."
      echo ""
      # If CentOS
      if [ "$OS" == "centos" ]; then
        echo "Removing yum versionlock for Salt."
        echo ""
        yum versionlock delete "salt-*"
        echo "Updating Salt packages and restarting services."
        echo ""
	if [ $is_airgap -eq 0 ]; then
          sh $UPDATE_DIR/salt/salt/scripts/bootstrap-salt.sh -r -F -M -x python3 stable "$NEWSALTVERSION"
	else 
          sh $UPDATE_DIR/salt/salt/scripts/bootstrap-salt.sh -F -M -x python3 stable "$NEWSALTVERSION"
	fi
        echo "Applying yum versionlock for Salt."
        echo ""
        yum versionlock add "salt-*"
      # Else do Ubuntu things
      elif [ "$OS" == "ubuntu" ]; then
        echo "Removing apt hold for Salt."
        echo ""
        apt-mark unhold "salt-common"
        apt-mark unhold "salt-master"
        apt-mark unhold "salt-minion"
        echo "Updating Salt packages and restarting services."
        echo ""
        sh $UPDATE_DIR/salt/salt/scripts/bootstrap-salt.sh -F -M -x python3 stable "$NEWSALTVERSION"
        echo "Applying apt hold for Salt."
        echo ""
        apt-mark hold "salt-common"
        apt-mark hold "salt-master"
        apt-mark hold "salt-minion"
      fi
}

verify_latest_update_script() {
    # Check to see if the update scripts match. If not run the new one.
    CURRENTSOUP=$(md5sum /opt/so/saltstack/default/salt/common/tools/sbin/soup | awk '{print $1}')
    GITSOUP=$(md5sum $UPDATE_DIR/salt/common/tools/sbin/soup | awk '{print $1}')
    if [[ "$CURRENTSOUP" == "$GITSOUP" ]]; then
      echo "This version of the soup script is up to date. Proceeding."
    else
      echo "You are not running the latest soup version. Updating soup."
      cp $UPDATE_DIR/salt/common/tools/sbin/soup $DEFAULT_SALT_DIR/salt/common/tools/sbin/
      salt-call state.apply common queue=True
      echo ""
      echo "soup has been updated. Please run soup again."
      exit 0
    fi
}

main () {
while getopts ":b" opt; do
  case "$opt" in
    b ) # process option b
       shift
       BATCHSIZE=$1
       if ! [[ "$BATCHSIZE" =~ ^[0-9]+$ ]]; then
         echo "Batch size must be a number greater than 0."
         exit 1
       fi
      ;;
    \? ) echo "Usage: cmd [-b]"
      ;;
  esac
done

echo "Checking to see if this is a manager."
echo ""
manager_check
echo "Checking to see if this is an airgap install"
echo ""
check_airgap
echo "Found that Security Onion $INSTALLEDVERSION is currently installed."
echo ""
detect_os
echo ""
if [ $is_airgap -eq 0 ]; then
  # Let's mount the ISO since this is airgap
  airgap_mounted
else
  echo "Cloning Security Onion github repo into $UPDATE_DIR."
  clone_to_tmp
fi
echo ""
echo "Verifying we have the latest soup script."
verify_latest_update_script
echo ""

echo "Let's see if we need to update Security Onion."
upgrade_check
space_check

echo "Checking for Salt Master and Minion updates."
upgrade_check_salt

echo ""
echo "Performing upgrade from Security Onion $INSTALLEDVERSION to Security Onion $NEWVERSION."
echo ""
echo "Updating dockers to $NEWVERSION."
update_dockers
echo ""
echo "Stopping Salt Minion service."
systemctl stop salt-minion
echo ""
echo "Stopping Salt Master service."
systemctl stop salt-master
echo ""

# Does salt need upgraded. If so update it.
if [ "$UPGRADESALT" == "1" ]; then
  echo "Upgrading Salt"
  # Update the repo files so it can actually upgrade
  if [ $is_airgap -eq 0 ]; then
    update_centos_repo
    yum clean all
  fi
  upgrade_salt
fi

echo "Making pillar changes."
pillar_changes
echo ""

# Only update the repo if its airgap
if [[ $is_airgap -eq 0 ]] && [[ "$UPGRADESALT" != "1" ]]; then
update_centos_repo
fi

echo ""
echo "Copying new Security Onion code from $UPDATE_DIR to $DEFAULT_SALT_DIR."
copy_new_files
echo ""
update_version

echo ""
echo "Locking down Salt Master for upgrade"
masterlock

echo ""
echo "Starting Salt Master service."
systemctl start salt-master

echo ""
echo "Running a highstate to complete the Security Onion upgrade on this manager. This could take several minutes."
highstate
echo ""
echo "Upgrade from $INSTALLEDVERSION to $NEWVERSION complete."

echo ""
echo "Stopping Salt Master to remove ACL"
systemctl stop salt-master

masterunlock

echo ""
echo "Starting Salt Master service."
systemctl start salt-master
highstate
playbook
unmount_update

if [ "$UPGRADESALT" == "1" ]; then
  echo ""
  echo "Upgrading Salt on the remaining Security Onion nodes from $INSTALLEDSALTVERSION to $NEWSALTVERSION."
  if [ $is_airgap -eq 0 ]; then
    salt -C 'not *_eval and not *_helix and not *_manager and not *_managersearch and not *_standalone' cmd.run "yum clean all"
  fi
  salt -C 'not *_eval and not *_helix and not *_manager and not *_managersearch and not *_standalone' -b $BATCHSIZE state.apply salt.minion 
  echo ""
fi

}

main "$@" | tee /dev/fd/3