#!/bin/bash

# this is a fix for the cve-2021-44228: https://www.randori.com/blog/cve-2021-44228/
# more info: https://www.lunasec.io/docs/blog/log4j-zero-day/

# what we do is delete org/apache/logging/log4j/core/lookup/JndiLookup.class from log4j-core-2*.jar files


if [[ $EUID -ne 0 ]]; then
  msg="Script must be run as root. Exiting"
  echo "$(date +"%M-%d-%Y %H:%M:%S.%N") [${prefix}]: ${log_level} ${msg}"
  exit 1
fi

PATCH_NAME="log4j-cve-2021-44228"
LOG_PREFIX="patch-${PATCH_NAME}"
RUN_ID=$(date +'%Y%m%d%H%M%S')
PATCH_COMPLETED_MARKER="/root/${PATCH_NAME}"
PATCH_ROOT_DIR="/root/${PATCH_NAME}_${RUN_ID}"
PATCH_TMP_DIR="/tmp/${PATCH_NAME}_${RUN_ID}"
LOG_FILE="${PATCH_ROOT_DIR}/log4j-cve-fix.log"
LOG4J_FILE_LIST_RAW="${PATCH_ROOT_DIR}/log4j2_jars_list_raw.txt"
LOG4J_FILE_LIST="${PATCH_ROOT_DIR}/log4j2_jars_list.txt"
STORAGE_FILE_LIST="${PATCH_ROOT_DIR}/storage_log4j2_jars_list.txt"
cluster_type=$(python -c "import hdinsight_common.ClusterManifestParser as ClusterManifestParser;cm=ClusterManifestParser.parse_local_manifest();print(cm.settings['cluster_type']);")
host_name=$(hostname -f)

if [ -f ${PATCH_COMPLETED_MARKER} ]; then
  echo "${PATCH_NAME} was already applied, last update: $(cat ${PATCH_COMPLETED_MARKER})" | logger -s -t ${LOG_PREFIX}
  exit 0
fi

mkdir -v -p ${PATCH_ROOT_DIR}
touch ${LOG_FILE}
touch ${LOG4J_FILE_LIST}
mkdir -v -p ${PATCH_TMP_DIR}
chmod -v -R 777 ${PATCH_TMP_DIR}

echo "$(date +"%M-%d-%Y %H:%M:%S.%N") [${LOG_PREFIX}]: INFO Starting to patch log4j2." >> ${LOG_FILE}

function execute_with_logs(){
  local command_to_execute=${*}

  logs=$(eval ${command_to_execute} 2>&1)
  exit_code=${?}
  log INFO "Output of running '${command_to_execute}':"
  OLD_IFS=${IFS}
  # changing field separator so that we can read line by line in for loop
  IFS=$'\n'
  for l in ${logs}
  do
     log INFO "${l}"
  done
  IFS=${OLD_IFS}
  return ${exit_code}
}

function log() {
  local prefix=${LOG_PREFIX}
  local log_level=${1:-INFO}
  local msg=${2:-"Empty Msg"}
  echo "$(date +"%M-%d-%Y %H:%M:%S.%N") [${prefix}]: ${log_level} ${msg}" >> ${LOG_FILE}
  if [[ "${log_level}" == "INFO" ]]; then
    logger -t ${prefix} -p user.info "${log_level} ${msg}"
  elif [[ "${log_level}" == "WARN" ]]; then
    logger -t ${prefix} -p user.warn "${log_level} ${msg}"
  elif [[ "${log_level}" == "ERROR" ]]; then
    logger -t ${prefix} -p user.err "${log_level} ${msg}"
  fi
}

function schedule_reboot() {
  ut=$(sudo awk '{print $1}' /proc/uptime | cut -d '.' -f 1)
  host_instance=$(sed 's/[^0-9]*\([0-9]*\).*/\1/' <<< $host_name)
  if [[ ( $host_name = hn* ) || ( $host_name = km* ) ]]; then
    # For headnodes and kafka management nodes reboot all nodes within 2 hours
    reboot_time="$((60 + (60 * ($host_instance % 2))))"
  elif [[ $host_name = wn* ]]; then
    # for worker nodes, reboot all nodes within 48 hour window
    reboot_time="$((10 + (30 * ($host_instance % 48))))"
  else
    msg="Not scheduling reboot, unknown node type. hostname=$host_name"
    log ERROR "$msg"
    return
  fi
  
  if [[ ${ut} -le 10800 ]]; then
    log INFO "Uptime was less than 3 hours, hence adding 180 mins to reboot time. up_time: ${ut}"
    reboot_time="$((180 + ${reboot_time}))"
  fi

  shutdown -r +$reboot_time "Rebooting to patch log4j cve patch"
  msg="Scheduled reboot of this node after $reboot_time minutes"
  log INFO "$msg"
  echo "$(date +"%M-%d-%Y %H:%M:%S.%N") ${msg}"
}

headnode_host=$(grep "headnodehost" /etc/hosts | awk '{print $2}' | tr -d ' ')

if [[ ${host_name} = ${headnode_host} ]];then
  log INFO "This is headnodehost. Patching jars in storage account."
  sudo -u hdfs hdfs --loglevel OFF dfs -find / -name "log4j-core-2*.jar" >> ${STORAGE_FILE_LIST}
  exit_code=$?
  log INFO "Fetching list of log4j-core-2 jars into file = $STORAGE_FILE_LIST ; exit_code= $exit_code"
  if [[ $exit_code -eq 0 ]]; then 
    echo "List of jars in storage account that need to be patched: $(cat ${STORAGE_FILE_LIST})" | logger -t ${LOG_PREFIX}

    for file_loc in $(cat ${STORAGE_FILE_LIST});
    do
      sudo rm -f -R ${PATCH_TMP_DIR}/*
      log INFO "Patching jar file in storage account = ${file_loc}"
      file_name="${PATCH_TMP_DIR}/$(basename ${file_loc})"
      execute_with_logs "sudo -u hdfs hdfs --loglevel OFF dfs -cp -f -ptpa ${file_loc} file:${PATCH_TMP_DIR}/"
      exit_code=$?
      log INFO "Downloaded from storage account file=$file_loc ; To= ${file_name} ; exit_code= $exit_code"
      sudo rm -f ${PATCH_TMP_DIR}/.*.crc
      execute_with_logs "zip -q -d ${file_name} org/apache/logging/log4j/core/lookup/JndiLookup.class"
      exit_code=$?
      log INFO "Patched jar file = $file_name ; exit_code= $exit_code"
      if [[ $exit_code -gt 0 ]]; then 
        continue
      fi
      execute_with_logs "chmod a+r ${file_name}"
      exit_code=$?
      log INFO "Updated permission to a+r for file = $file_name ; exit_code= $exit_code"
      if [[ $exit_code -gt 0 ]]; then 
        continue
      fi
      execute_with_logs "sudo -u hdfs hdfs --loglevel OFF dfs -cp -f -ppa file:${file_name} ${file_loc}"
      exit_code=$?
      log INFO "Uploaded jar file = $file_name ; To= ${file_loc} ; exit_code= $exit_code"
    done
  fi
else
  log INFO "This is not headnodehost. Skipping patching of jars in storage account"
fi

# execute find and for each file found,
#  a. collect the jar locations,
#  b. back up the jar file

echo "Collecting jars that need to be patched" | logger -t ${LOG_PREFIX}

if [[ $host_name = ed* ]]; then
  log INFO "This node is an edgenode and patching will done to hdiinsight specific jars and will not reboot node. HostName= $host_name."
  sudo find /usr/hdp/ /usr/lib/ambari-infra-solr/ /var/lib/ambari-server/ -type f -name "log4j-core-2*.jar" -not -path ${PATCH_ROOT_DIR} -exec sh -c "jar tvf {} | grep -q -i JndiLookup.class && echo {} >> ${LOG4J_FILE_LIST}" \;
else
  log INFO "This node is NOT an edgenode and patching will be done on all jars in known locations. HostName= $host_name."
  sudo find /var/ /opt/ /usr/ -type f -name "log4j-core-2*.jar" -not -path ${PATCH_ROOT_DIR} -exec sh -c "jar tvf {} | grep -q -i JndiLookup.class && echo {} >> ${LOG4J_FILE_LIST}" \;
fi

# find shaded jars
sudo find /usr/hdp/ -type f -name "spark-examples-*.jar" -not -path ${PATCH_ROOT_DIR} -exec sh -c "jar tvf {} | grep -q -i JndiLookup.class && echo {} >> ${LOG4J_FILE_LIST}" \;
sudo find /usr/hdp/ -type f -name "hive-jdbc-*-standalone.jar" -not -path ${PATCH_ROOT_DIR} -exec sh -c "jar tvf {} | grep -q -i JndiLookup.class && echo {} >> ${LOG4J_FILE_LIST}" \;
sudo find /usr/lib/hdinsight-kafka-restproxy-agent -type f -name "hdinsight-kafka-restproxy.jar" -not -path ${PATCH_ROOT_DIR} -exec sh -c "jar tvf {} | grep -q -i JndiLookup.class && echo {} >> ${LOG4J_FILE_LIST}" \;

# log all positive files found
echo "List of jars that need to be patched: $(cat ${LOG4J_FILE_LIST})" | logger -t ${LOG_PREFIX}

for file_loc in $(cat ${LOG4J_FILE_LIST});
do
  log INFO "Patching jar file = $file_loc"
  execute_with_logs "cp --backup=numbered $file_loc ${PATCH_ROOT_DIR}"
  execute_with_logs "zip -q -d $file_loc org/apache/logging/log4j/core/lookup/JndiLookup.class"
  exit_code=$?
  log INFO "Patched jar file = $file_loc ; exit_code= $exit_code"
done

if [[ ( $host_name = hn* ) || ( $host_name = km* ) ]]; then
  # Reboot headnodes and Kafka rest proxy nodes for sure on all cluster types
  schedule_reboot
else
  # all other nodes do not reboot, like for eg. zookeeper nodes
  log INFO "Not scheduling reboot of this node."
fi

msg="Script to remove JndiLookup.class from log4j-core-2* is completed."
log INFO "$msg"
echo "$(date +"%M-%d-%Y %H:%M:%S.%N") ${msg}" | logger -s -t ${LOG_PREFIX} | tee ${PATCH_COMPLETED_MARKER}