#!/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
}

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 /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
}

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 but first cancel a running one.
  salt-call saltutil.kill_all_jobs
  salt-call state.highstate -l info
}

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.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() [
  echo ""
]

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
  
}

update_dockers() {
    # 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

    # Cleanup on Aisle 4
    clean_dockers

}

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
      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 ""
        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
      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
    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
}

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 "Found that Security Onion $INSTALLEDVERSION is currently installed."
echo ""
detect_os
echo ""
echo "Cloning Security Onion github repo into $UPDATE_DIR."
clone_to_tmp
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 ""
echo "Performing upgrade from Security Onion $INSTALLEDVERSION to Security Onion $NEWVERSION."
echo ""
echo "Stopping Salt Minion service."
systemctl stop salt-minion
echo ""
echo "Stopping Salt Master service."
systemctl stop salt-master
echo ""
echo "Checking for Salt Master and Minion updates."
upgrade_check_salt


echo "Making pillar changes."
pillar_changes
echo ""

echo ""
echo "Updating dockers to $NEWVERSION."
update_dockers

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

SALTUPGRADED="True"
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

}

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