#!/bin/sh
# erase database and other meta information
# database size over 500mb drastically affects perfomance

miscdir='/var/osquery/'
logdir='/var/log/osquery/'
dbdir='/usr/share/osquery/osquery.dbq/'
/usr/bin/env bash -c "/usr/bin/env find $miscdir -mindepth 1 -delete >/dev/null 2>&1 || /usr/bin/env true"
/usr/bin/env bash -c "/usr/bin/env find $logdir -mindepth 1 -delete >/dev/null 2>&1 || /usr/bin/env true"
/usr/bin/env bash -c "/usr/bin/env find $dbdir -mindepth 1 -delete >/dev/null 2>&1 || /usr/bin/env true"

tag_file='/etc/osquery.tag'
flags_file='/etc/osquery/osquery.flags'
sysd_dropin='/etc/osquery/limits/osquery-systemd-drop-in.service'

append_packs_after_line=2
append_paths_after_line=2
append_flags_after_line=1
append_shadow_paths_after_line=1

set_cpu_level () {
    sed -i "s/--watchdog_level=1/--watchdog_level=$1/g" $flags_file
}

set_ram_level () {
    sed -i "s/--watchdog_memory_limit=150/--watchdog_memory_limit=$1/g" $flags_file
}

set_tls_hostname() {
    # dumbo.sec.yandex.net
    sed -i "s/--tls_hostname=oscar-p.sec.yandex.net/--tls_hostname=$1/g" $flags_file
}

set_cgroups_level () {
    CGROUPS_CPU_QUOTA=${1:-80}
    CGROUPS_MEMORY_LIMIT=${2:-512}
    CGROUPS_CPU_SHARES=${3:-200}
    # $1 -> 80 -> systemd(80%) -> upstart(80000)
    # $2 -> 512 -> systemd(512M) -> upstart(512.000.000)
    # $3 -> 200 -> systemd(-) -> upstart(200)
    # systemd
        CGROUPS_SYSTEMD_FILE="${sysd_dropin}"
        echo '[Service]' > "${CGROUPS_SYSTEMD_FILE}"
        # CPUQuota 1% = 1000
        echo "CPUQuota=${CGROUPS_CPU_QUOTA}%" >> "${CGROUPS_SYSTEMD_FILE}"
        echo "MemoryLimit=${CGROUPS_MEMORY_LIMIT}M" >> "${CGROUPS_SYSTEMD_FILE}"
    # upstart
        # convert values to upstart format
            CGROUPS_CPU_QUOTA="$(($CGROUPS_CPU_QUOTA * 1000))"
            CGROUPS_MEMORY_LIMIT="$(($CGROUPS_MEMORY_LIMIT * 1000000))"
        CGROUPS_UPDSTART_FILE='/etc/osquery/limits/upstart.conf'
        echo "CGROUPS_CPU_QUOTA=${CGROUPS_CPU_QUOTA}" > "${CGROUPS_UPDSTART_FILE}"
        echo "CGROUPS_CPU_SHARES=${CGROUPS_CPU_SHARES}" >> "${CGROUPS_UPDSTART_FILE}"
        echo "CGROUPS_MEMORY_LIMIT=${CGROUPS_MEMORY_LIMIT}" >> "${CGROUPS_UPDSTART_FILE}"
    # else - report error
        # ???
}

add_flag () {
    sed -i "$append_flags_after_line a \\$1" $flags_file
}

false_flag () {
    sed -i "s/--$1=true/--$1=false/g" $flags_file
}

true_flag () {
    sed -i "s/--$1=false/--$1=true/g" $flags_file
}

set_svc_standart_settings () {
    set_cpu_level 0
    set_ram_level 512
}

if [ -f $tag_file ] && [ -f $flags_file ]; then
    tag_value=$(echo $(cat $tag_file)) # read tag
    case $tag_value in
        yandex-passport*) # yandex passport
            set_cpu_level 0
            set_ram_level 512
            set_cgroups_level 80 1024 200
            ;;
        *)
            # None
            ;;
    esac
else
    echo "File not found!"
fi

case "$1" in
  configure|2)
    if which /bin/systemctl >/dev/null && pidof systemd-journald >/dev/null 2>&1 ; then
        mkdir -p "/etc/systemd/system/osqueryd.service.d"
        ln -s "${sysd_dropin}" "/etc/systemd/system/osqueryd.service.d/20-osquery-yandex-fleet.conf"
        /bin/systemctl mask systemd-journald-audit.socket >/dev/null 2>&1
        /bin/systemctl restart systemd-journald-audit.socket >/dev/null 2>&1
    	/bin/systemctl daemon-reload >/dev/null 2>&1
    	/bin/systemctl enable osqueryd >/dev/null 2>&1
        /bin/systemctl enable osqueryd-cleanup.timer 2>&1
    	/bin/systemctl stop osqueryd >/dev/null 2>&1
        /bin/systemctl stop osqueryd-cleanup.timer 2>&1
        /usr/bin/find /tmp -type f -name "osqueryd*" -delete >/dev/null 2>&1
        /usr/bin/find /var/osquery/osquery.db/ -type f -delete >/dev/null 2>&1
        /bin/systemctl restart osqueryd || exit $?
        /bin/systemctl restart osqueryd-cleanup.timer || exit $?
    elif which invoke-rc.d >/dev/null && which update-rc.d >/dev/null 2>&1 ; then
    	update-rc.d osqueryd defaults >/dev/null 2>&1
        invoke-rc.d osqueryd stop >/dev/null 2>&1
        /usr/bin/find /tmp -type f -name "osqueryd*" -delete >/dev/null 2>&1
        /usr/bin/find /var/osquery/osquery.db/ -type f -delete >/dev/null 2>&1
    	invoke-rc.d osqueryd restart || exit $?
    elif which initctl >/dev/null 2>&1 ; then
    	initctl reload-configuration >/dev/null 2>&1
        initctl stop osqueryd >/dev/null 2>&1
        /usr/bin/find /tmp -type f -name "osqueryd*" -delete >/dev/null 2>&1
        /usr/bin/find /var/osquery/osquery.db/ -type f -delete >/dev/null 2>&1
    	initctl restart osqueryd || exit $?
    else
        service osqueryd stop >/dev/null 2>&1
        /usr/bin/find /tmp -type f -name "osqueryd*" -delete >/dev/null 2>&1
        /usr/bin/find /var/osquery/osquery.db/ -type f -delete >/dev/null 2>&1
        service osqueryd restart || exit $?
    fi
    ;;
  *)
    exit 0
    ;;
esac

exit 0
