#!/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
SOUP_BRANCH=$1
BATCHSIZE=5

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')$ ]]; then
    echo "This is a manager. We can proceed"
  else
    echo "Please run soup on the manager. The manager controls all updates."
    exit 0
  fi
}

clean_dockers() {
  # Place Holder for cleaning up old docker images
  echo ""
}

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
  if [ -n "$SOUP_BRANCH" ]; then
    git clone -b $SOUP_BRANCH https://github.com/Security-Onion-Solutions/securityonion.git
  else
    git clone https://github.com/Security-Onion-Solutions/securityonion.git
  fi
  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 /tmp/sogh/securityonion
  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
}

highstate() {
  # Run a highstate but first cancel a running one.
  salt-call saltutil.kill_all_jobs
  salt-call state.highstate -l info
}

pillar_changes() {
    # This function is to add any new pillar items if needed.
    echo "Checking to see if pillar changes are needed"

}

update_dockers() {
    # List all the containers
    if [ $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-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

}

update_version() {
  # Update the version to the latest
  echo "Updating the version file."
  echo $NEWVERSION > /etc/soversion
  sed -i "s/$INSTALLEDVERSION/$NEWVERSION/g" /opt/so/saltstack/local/pillar/static.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
      SALTUPGRADED=True
      echo "Performing upgrade of Salt from $INSTALLEDSALTVERSION to $NEWSALTVERSION"
      echo ""
      # If CentOS
      echo "Removing yum versionlock for Salt"
      echo ""
      yum versionlock delete "salt-*"
      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 yum versionlock for Salt"
      echo ""
      yum versionlock add "salt-*"
      # Else do Ubuntu things
    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 /tmp/sogh/securityonion/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
}

while getopts ":b" opt; do
  case ${opt} in
    b ) # process option b
      if [[ $OPTARG =~ ^?[0-9]+$ ]] && [[ $OPTARG -gt 0 ]]; then
        BATCHSIZE=$OPTARG
      else
        echo "Batch size must be a number greater than 0"
      fi
      ;;
    \? ) echo "Usage: cmd [-b]"
      ;;
  esac
done

echo "Checking to see if this is a manager"
manager_check
echo "Cloning latest code to a temporary location"
clone_to_tmp
echo ""
echo "Verifying we have the latest script"
verify_latest_update_script
echo ""

echo "Let's see if we need to update"
upgrade_check


echo ""
echo "Performing Upgrade from $INSTALLEDVERSION to $NEWVERSION"
echo ""
echo "Stopping Salt Master service"
systemctl stop salt-master
echo ""
echo "Stopping Salt Minion service"
systemctl stop salt-minion
echo ""
echo "Checking for Salt updates"
upgrade_check_salt


echo "Making pillar changes"
pillar_changes
echo ""

echo "Cleaning up old dockers"
clean_dockers
echo ""
echo "Updating docker to $NEWVERSION"
update_dockers

echo ""
echo "Copying new code"
copy_new_files
echo ""
echo "Updating version"
update_version


echo ""
echo "Starting Salt Master service"
systemctl start salt-master
#echo ""
#echo "Starting Salt Minion service"
#systemctl start salt-minion
#echo ""


echo ""
echo "Running a highstate to complete upgrade"
highstate
echo ""
echo "Upgrade from $INSTALLEDVERSION to $NEWVERSION complete."

if [ "$SALTUPGRADED" = True ]; then
  echo ""
  echo "Upgrading Salt on the remaining Security Onion nodes from $INSTALLEDSALTVERSION to $NEWSALTVERSION."
  salt -C 'not *_eval and not *_helix and not *_manager and not *_managersearch and not *_standalone' -b $BATCHSIZE state.apply salt.minion 
  echo ""
fi