#!/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)
default_salt_dir=/opt/so/saltstack/default

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

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

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" \
      "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/$NEWVERISON/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
    else
      echo "Performing Upgrade from $INSTALLEDVERSION to $NEWVERSION"
    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
}

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 "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 "Running a highstate to complete upgrade"
highstate
echo ""
echo "Updating version"
update_version
echo ""
echo "Upgrade from $INSTALLEDVERSION to $NEWVERSION complete."