#!/bin/sh
#
# osqueryd  Start/Stop the osquery daemon.
#
# chkconfig: 345 90 60
# Description:
#  Ya osquery vanilla build
#
### BEGIN INIT INFO
# Provides: osquery osqueryd
# Required-Start: $local_fs $syslog
# Required-Stop: $local_fs $syslog
# Default-Start:  3 4 5
# Default-Stop: 0 1 6
# Short-Description: run osqueryd daemon
# Description:
#   Ya osquery vanilla build
#
### END INIT INFO

if [ -z $RETVAL ]; then RETVAL=0; fi
if [ -z $PROG ]; then PROG="osqueryd"; fi
if [ -z $EXEC ]; then EXEC=/usr/bin/osqueryd; fi
if [ -z $FLAGS_PATH ]; then FLAGS_PATH=/etc/osquery/osquery.flags; fi
if [ -z $REAL_CONFIG_PATH ]; then REAL_CONFIG_PATH=/etc/osquery/osquery.conf; fi
if [ -z $LOCKFILE ]; then LOCKFILE=/var/lock/osqueryd; fi
if [ -z $PIDFILE ]; then PIDFILE=/var/run/osqueryd.pidfile; fi
if [ -z $OLD_PIDFILE ]; then OLD_PIDFILE=/var/run/osqueryd.pid; fi
if [ -z $UID ]; then UID=$(id -u); fi

if [ $UID -eq 0 ] && [ -e /etc/default/$PROG ]; then
  . /etc/default/$PROG
fi

if [ -e /etc/init.d/functions ]; then
  . /etc/init.d/functions
fi

if [ ! -e $FLAGS_PATH ] && [ ! -e $REAL_CONFIG_PATH ]; then
  echo "No config file found at $REAL_CONFIG_PATH"
  echo "Additionally, no flags file or config override found at $FLAGS_PATH"
  echo "See '/usr/share/osquery/osquery.example.conf' for an example config."
  RETVAL=1
fi

move_pidfile() {
  if [ -f $OLD_PIDFILE ]; then
    # Support for deprecated pidfile location.
    mv $OLD_PIDFILE $PIDFILE
  fi
}

ensure_root() {
  if [ $UID -ne 0 ] ; then
    echo "User has insufficient privilege."
    RETVAL=1
  fi
}

clean_database() {

  # /usr/bin/env bash -c '/usr/bin/env find /var/osquery/ -mindepth 1 -delete || /usr/bin/env true'
  # /usr/bin/env bash -c '/usr/bin/env find /var/log/osquery/ -mindepth 1 -delete || /usr/bin/env true'
  # /usr/bin/env bash -c '/usr/bin/env find /usr/share/osquery/osquery.dbq/ -mindepth 1 -delete || /usr/bin/env true'

  /usr/bin/env find /usr/share/osquery/osquery.dbq/ -mindepth 1 -mtime +0 -type f \! \( -name "CURRENT" -o -name "IDENTITY" -o -name "LOCK" -o -name "MANIFEST-*" \) -delete > /dev/null 2>&1 || /usr/bin/env true
  /usr/bin/env find /var/osquery/ -mindepth 1 -mtime +0 -type f \! \( -name "CURRENT" -o -name "IDENTITY" -o -name "LOCK" -o -name "MANIFEST-*" \) -delete > /dev/null 2>&1 || /usr/bin/env true
  /usr/bin/env find /var/log/osquery/ -mindepth 1 -mtime +0 -type f \! \( -name "CURRENT" -o -name "IDENTITY" -o -name "LOCK" -o -name "MANIFEST-*" \) -delete > /dev/null 2>&1 || /usr/bin/env true

  if [ $(du -k /usr/share/osquery/osquery.dbq/ | cut -f1) -ge 350000 ]; then
      { rm /usr/share/osquery/osquery.dbq/* && service osqueryd restart; } || exit 255
  fi
}

add_to_crontab() {
  # croncmd="service osqueryd restart 2>&1"
  # cronjob="0 */84 * * * $croncmd"
  # ( crontab -l | grep -v -F "$croncmd" ; echo "$cronjob" ) | crontab -

CRONTABSTRING='0 */84 * * * root service osqueryd restart 2>&1'
if [ -d /etc/cron.d/ ]; then
  if [ -f /etc/cron.d/yandex-osquery ]; then
    if ! grep -q "$CRONTABSTRING" /etc/cron.d/yandex-osquery; then
      echo "$CRONTABSTRING" > /etc/cron.d/yandex-osquery
      sudo service cron restart
    fi
  else
    mkdir -p /etc/cron.d/
    echo "$CRONTABSTRING" > /etc/cron.d/yandex-osquery
    sudo service cron restart
  fi
fi
}

mount_tmpfs() {
  test -x /etc/osquery/osqueryd-tmpfs.sh && /etc/osquery/osqueryd-tmpfs.sh mount
}

umount_tmpfs() {
  test -x /etc/osquery/osqueryd-tmpfs.sh && /etc/osquery/osqueryd-tmpfs.sh umount
}

add_to_cgroup() {
  # CGROUPS 
  CGROUPS_MEMORY_LIMIT=536870912
  CGROUPS_CPU_SHARES=200
  # CPUQuota 1% = 1000
  CGROUPS_CPU_QUOTA=80000

  if [ -f "/etc/osquery/limits/upstart.conf" ]; then
    . /etc/osquery/limits/upstart.conf
  fi
  
  for x in /etc/osquery/limits/upstart.conf.d/*.conf ; do
    if [ -e "${x}" ]; then
      . "${x}"
    fi
  done

  if ! [ -d "/sys/fs/cgroup/cpu,cpuacct" ]; then
    mkdir -p /sys/fs/cgroup/cpu,cpuacct
  fi
  if ! mountpoint -q /sys/fs/cgroup/cpu,cpuacct; then
    mount -t cgroup -o cpu,cpuacct none /sys/fs/cgroup/cpu,cpuacct > /dev/null 2>&1
  fi
  if ! [ -d "$/sys/fs/cgroup/cpu,cpuacct/osquery" ]; then
    mkdir -p /sys/fs/cgroup/cpu,cpuacct/osquery > /dev/null 2>&1
  fi
  echo $CGROUPS_CPU_SHARES 2>/dev/null > /sys/fs/cgroup/cpu,cpuacct/osquery/cpu.shares
  echo $CGROUPS_CPU_QUOTA 2>/dev/null > /sys/fs/cgroup/cpu,cpuacct/osquery/cpu.cfs_quota_us
  echo "$(cat $PIDFILE)" 2>/dev/null > /sys/fs/cgroup/cpu,cpuacct/osquery/cgroup.procs
  echo -n "$(pgrep -n osqueryd)" 2>/dev/null >> /sys/fs/cgroup/cpu,cpuacct/osquery/cgroup.procs

  # for 12 ubuntu
  if grep -i "DISTRIB_RELEASE=12" /etc/lsb-release > /dev/null 2>&1; then
    if ! [ -d "$/sys/fs/cgroup/cpu" ]; then
      mkdir -p /sys/fs/cgroup/cpu > /dev/null 2>&1
    fi
    if ! mountpoint -q /sys/fs/cgroup/cpu; then
      mount -t cgroup -o cpu none /sys/fs/cgroup/cpu > /dev/null 2>&1
    fi
    if ! [ -d "$/sys/fs/cgroup/cpu,cpuacct/osquery" ]; then
      mkdir -p /sys/fs/cgroup/cpu/osquery
    fi
  echo $CGROUPS_CPU_SHARES 2>/dev/null > /sys/fs/cgroup/cpu/osquery/cpu.shares
  echo $CGROUPS_CPU_QUOTA 2>/dev/null > /sys/fs/cgroup/cpu/osquery/cpu.cfs_quota_us
  echo "$(cat $PIDFILE)" 2>/dev/null > /sys/fs/cgroup/cpu/osquery/cgroup.procs
  echo -n "$(pgrep -n osqueryd)" 2>/dev/null >> /sys/fs/cgroup/cpu/osquery/cgroup.procs
  fi

  if ! [ -d "$/sys/fs/cgroup/memory" ]; then
    mkdir -p /sys/fs/cgroup/memory > /dev/null
  fi
  if ! mountpoint -q /sys/fs/cgroup/memory; then
    mount -t cgroup -o memory none /sys/fs/cgroup/memory  > /dev/null
  fi
  if ! [ -d "$/sys/fs/cgroup/memory/osquery" ]; then
    mkdir -p /sys/fs/cgroup/memory/osquery
  fi
  echo $CGROUPS_MEMORY_LIMIT 2>/dev/null > /sys/fs/cgroup/memory/osquery/memory.limit_in_bytes
  echo "$(cat $PIDFILE)" 2>/dev/null > /sys/fs/cgroup/memory/osquery/cgroup.procs
  echo -n "$(pgrep -n osqueryd)" 2>/dev/null >> /sys/fs/cgroup/memory/osquery/cgroup.procs
}

start() {
  ensure_root
  move_pidfile
  clean_database
  add_to_crontab
  mount_tmpfs

  ARGS=""
  if [ -f $PIDFILE ]; then
    PID=$(cat $PIDFILE)
    PROCNAME=$(ps -p $PID -o comm\=)
    if [ "$PROCNAME" = "$PROG" ]; then
      return 0
    else
      # osqueryd pidfile exists but it's not running
      rm $PIDFILE
    fi
  fi

  if [ -e $FLAGS_PATH ]; then ARGS="$ARGS --flagfile=$FLAGS_PATH"; fi
  if [ -e $REAL_CONFIG_PATH ]; then ARGS="$ARGS --config_path=$REAL_CONFIG_PATH"; fi

  $EXEC $ARGS \
        --pidfile=$PIDFILE \
        --daemonize=true
  RETVAL=$?
  
  CONFIG_FILE="/etc/default/osqueryd.conf"
  if [ -f "${CONFIG_FILE}" ]; then
    . "${CONFIG_FILE}"
    if [ ! -z "${ENABLE_CGROUPS}" ]; then
      if [ "$ENABLE_CGROUPS" = 'false' ]; then
        :
      else
        add_to_cgroup
      fi
    fi
  else
    add_to_cgroup
  fi
}

stop() {
  ensure_root
  move_pidfile

  if [ ! -f $PIDFILE ] ; then
    RETVAL=0
  else
    PID=$(cat $PIDFILE)
    # Terminate the daemon and watchers
    pkill -g $PID
    # Allow the event threads to tear down
    ( while kill -0 $PID >/dev/null 2>&1; do sleep 0.2; done ) & DPID=$!
    # If the event threads are still running after 5 seconds, kill them
    ( sleep 5 && pkill -9 -g $PID && kill -9 $DPID ) 2>/dev/null & WPID=$!
    if wait $DPID 2>/dev/null; then
      pkill -9 -P $WPID
      wait $WPID
    fi
    rm -f $PIDFILE
  fi

  umount_tmpfs
}

restart() {
  stop
  start
}

status() {
  if [ -f $PIDFILE ]; then
    PID=$(cat $PIDFILE)
    PROCNAME=$(ps -p $PID -o comm\=)
    if [ "$PROCNAME" = "$PROG" ]; then
      echo "$PROG is already running: $PID"
      RETVAL=0
    else
      # osqueryd pidfile exists but it's not running
      echo "$PROG is not running but a stale pidfile was found."
      RETVAL=7
    fi
  elif [ -f $OLD_PIDFILE ]; then
    PID=$(cat $OLD_PIDFILE)
    PROCNAME=$(ps -p $PID -o comm\=)
    if [ "$PROCNAME" = "$PROG" ]; then
      echo "$PROG is already running (old pidfile): $PID"
      RETVAL=0
    else
      # osqueryd pidfile exists but it's not running
      echo "$PROG is not running but a stale old pidfile was found."
      RETVAL=7
    fi
  else
    echo "$PROG is not running. no pidfile found."
    RETVAL=7
  fi
}

usage() {
  echo "Usage: $0 {start|stop|status|restart}"
  RETVAL=2
}

case "$1" in
    start) start ;;
    stop) stop ;;
    restart) restart ;;
    status) status ;;
    *) usage ;;
esac

exit ${RETVAL}
