#!/bin/bash
#
# Simple script that desperately tries to load sysdig-probe looking
# for it in a bunch of ways. Convenient when running sysdig inside
# a container or in other weird environments.
#

get_kernel_config() {
	if [ -f /proc/config.gz ]; then
		echo "Found kernel config at /proc/config.gz"
		HASH=$(zcat /proc/config.gz | md5sum - | cut -d' ' -f1)
	elif [ -f "/boot/config-${KERNEL_RELEASE}" ]; then
		echo "Found kernel config at /boot/config-${KERNEL_RELEASE}"
		HASH=$(md5sum "/boot/config-${KERNEL_RELEASE}" | cut -d' ' -f1)
	elif [ ! -z "${SYSDIG_HOST_ROOT}" ] && [ -f "${SYSDIG_HOST_ROOT}/boot/config-${KERNEL_RELEASE}" ]; then
		echo "Found kernel config at ${SYSDIG_HOST_ROOT}/boot/config-${KERNEL_RELEASE}"
		HASH=$(md5sum "${SYSDIG_HOST_ROOT}/boot/config-${KERNEL_RELEASE}" | cut -d' ' -f1)
	elif [ -f "/usr/lib/ostree-boot/config-${KERNEL_RELEASE}" ]; then
		echo "Found kernel config at /usr/lib/ostree-boot/config-${KERNEL_RELEASE}"
		HASH=$(md5sum "/usr/lib/ostree-boot/config-${KERNEL_RELEASE}" | cut -d' ' -f1)
	elif [ ! -z "${SYSDIG_HOST_ROOT}" ] && [ -f "${SYSDIG_HOST_ROOT}/usr/lib/ostree-boot/config-${KERNEL_RELEASE}" ]; then
		echo "Found kernel config at ${SYSDIG_HOST_ROOT}/usr/lib/ostree-boot/config-${KERNEL_RELEASE}"
		HASH=$(md5sum "${SYSDIG_HOST_ROOT}/usr/lib/ostree-boot/config-${KERNEL_RELEASE}" | cut -d' ' -f1)
	elif [ -f /lib/modules/${KERNEL_RELEASE}/config ]; then
		# this code works both for native host and agent container assuming that
		# Dockerfile sets up the desired symlink /lib/modules -> $SYSDIG_HOST_ROOT/lib/modules
		echo "Found kernel config at /lib/modules/${KERNEL_RELEASE}/config"
		HASH=$(md5sum "/lib/modules/${KERNEL_RELEASE}/config" | cut -d' ' -f1)
	fi

	if [ -z "${HASH}" ]; then
		echo "Cannot find kernel config"
		exit 1
	fi
}

load_kernel_probe() {
	if ! hash lsmod > /dev/null 2>&1; then
		echo "This program requires lsmod"
		exit 1
	fi

	if ! hash modprobe > /dev/null 2>&1; then
		echo "This program requires modprobe"
		exit 1
	fi

	if ! hash rmmod > /dev/null 2>&1; then
		echo "This program requires rmmod"
		exit 1
	fi

	echo "* Unloading ${PROBE_NAME}, if present"
	rmmod "${PROBE_NAME}"

	if lsmod | grep "$(echo "${PROBE_NAME}" | tr "-" "_")" > /dev/null 2>&1; then
		echo "* ${PROBE_NAME} seems to still be loaded, hoping the best"
		exit 0
	fi

	# skip dkms on UEK hosts because it will always fail
	if [[ $(uname -r) == *uek* ]]; then
		echo "* Skipping dkms install for UEK host"
	else
		echo "* Running dkms install for ${PACKAGE_NAME}"
		if dkms install -m "${PACKAGE_NAME}" -v "${SYSDIG_VERSION}" -k "${KERNEL_RELEASE}"; then
			echo "* Trying to load a dkms ${PROBE_NAME}, if present"

			if insmod "/var/lib/dkms/${PACKAGE_NAME}/${SYSDIG_VERSION}/${KERNEL_RELEASE}/${ARCH}/module/${PROBE_NAME}.ko" > /dev/null 2>&1; then
				echo "${PROBE_NAME} found and loaded in dkms"
				exit 0
			elif insmod "/var/lib/dkms/${PACKAGE_NAME}/${SYSDIG_VERSION}/${KERNEL_RELEASE}/${ARCH}/module/${PROBE_NAME}.ko.xz" > /dev/null 2>&1; then
				echo "${PROBE_NAME} found and loaded in dkms (xz)"
				exit 0
			else
				echo "* Unable to insmod"
			fi
		else
			DKMS_LOG="/var/lib/dkms/${PACKAGE_NAME}/${SYSDIG_VERSION}/build/make.log"
			if [ -f "${DKMS_LOG}" ]; then
				echo "* Running dkms build failed, dumping ${DKMS_LOG}"
				cat "${DKMS_LOG}"
			else
				echo "* Running dkms build failed, couldn't find ${DKMS_LOG}"
			fi
		fi
	fi

	echo "* Trying to load a system ${PROBE_NAME}, if present"

	if modprobe "${PROBE_NAME}" > /dev/null 2>&1; then
		echo "${PROBE_NAME} found and loaded with modprobe"
		exit 0
	fi

	echo "* Trying to find precompiled ${PROBE_NAME} for ${KERNEL_RELEASE}"

	get_kernel_config

	local SYSDIG_PROBE_FILENAME="${PROBE_NAME}-${SYSDIG_VERSION}-${ARCH}-${KERNEL_RELEASE}-${HASH}.ko"

	if [ -f "${HOME}/.sysdig/${SYSDIG_PROBE_FILENAME}" ]; then
		echo "Found precompiled module at ~/.sysdig/${SYSDIG_PROBE_FILENAME}, loading module"
		insmod "${HOME}/.sysdig/${SYSDIG_PROBE_FILENAME}"
		exit $?
	fi

	local URL
	URL=$(echo "${SYSDIG_PROBE_URL}/${SYSDIG_REPOSITORY}/sysdig-probe-binaries/${SYSDIG_PROBE_FILENAME}" | sed s/+/%2B/g)

	echo "* Trying to download precompiled module from ${URL}"
	if curl --create-dirs -f -s -o "${HOME}/.sysdig/${SYSDIG_PROBE_FILENAME}" "${URL}"; then
		echo "Download succeeded, loading module"
		insmod "${HOME}/.sysdig/${SYSDIG_PROBE_FILENAME}"
		exit $?
	else
		echo "Download failed, consider compiling your own ${PROBE_NAME} and loading it or getting in touch with the sysdig community"
		exit 1
	fi
}

load_bpf_probe() {
	echo "* Mounting debugfs"

	if [ ! -d /sys/kernel/debug/tracing ]; then
		mount -t debugfs nodev /sys/kernel/debug
	fi

	get_kernel_config

	if [ ! -z "${SYSDIG_HOST_ROOT}" ] && [ -f "${SYSDIG_HOST_ROOT}/etc/os-release" ]; then
		. "${SYSDIG_HOST_ROOT}/etc/os-release"

		if [ "${ID}" == "cos" ]; then
			COS=1
		fi
	fi

	local BPF_PROBE_FILENAME="${BPF_PROBE_NAME}-${SYSDIG_VERSION}-${ARCH}-${KERNEL_RELEASE}-${HASH}.o"

	if [ ! -f "${HOME}/.sysdig/${BPF_PROBE_FILENAME}" ]; then
		if [ -n "${COS}" ]; then
			echo "* COS detected (build ${BUILD_ID}), downloading and setting up kernel headers"

			local -r download_url="https://storage.googleapis.com/cos-tools/${BUILD_ID}/kernel-src.tar.gz"

			echo "* Downloading ${download_url}"

			mkdir -p /tmp/kernel
			cd /tmp/kernel
			if ! curl --create-dirs -s -S -f -O "${download_url}"; then
				exit 1;
			fi

			echo "* Extracting kernel sources"

			tar xf kernel-src.tar.gz
			zcat /proc/config.gz > .config
			sed -i 's/LOCALVERSION=""/LOCALVERSION="+"/' .config

			echo "* Configuring kernel"

			make olddefconfig > /dev/null
			make modules_prepare > /dev/null

			export KERNELDIR=/tmp/kernel
		fi

		echo "* Trying to compile BPF probe ${BPF_PROBE_NAME} (${BPF_PROBE_FILENAME})"

		make -C "/usr/src/${PACKAGE_NAME}-${SYSDIG_VERSION}/bpf" > /dev/null

		mkdir -p ~/.sysdig
		mv "/usr/src/${PACKAGE_NAME}-${SYSDIG_VERSION}/bpf/probe.o" "${HOME}/.sysdig/${BPF_PROBE_FILENAME}"

		if [ -n "${COS}" ]; then
			rm -r /tmp/kernel
		fi
	fi

	if [ ! -f "${HOME}/.sysdig/${BPF_PROBE_FILENAME}" ]; then
		local URL
		URL=$(echo "${SYSDIG_PROBE_URL}/${SYSDIG_REPOSITORY}/sysdig-probe-binaries/${BPF_PROBE_FILENAME}" | sed s/+/%2B/g)

		echo "* Trying to download precompiled BPF probe from ${URL}"

		curl --create-dirs -f -s -S -o "${HOME}/.sysdig/${BPF_PROBE_FILENAME}" "${URL}"
	fi

	if [ -f "${HOME}/.sysdig/${BPF_PROBE_FILENAME}" ]; then
		if [ ! -f /proc/sys/net/core/bpf_jit_enable ]; then
			echo "**********************************************************"
			echo "** BPF doesn't have JIT enabled, performance might be   **"
			echo "** degraded. Please ensure to run on a kernel with      **"
			echo "** CONFIG_BPF_JIT enabled and/or use --net=host if      **"
			echo "** running inside a container.                          **"
			echo "**********************************************************"
		fi

		echo "* BPF probe located, it's now possible to start sysdig"

		ln -sf "${HOME}/.sysdig/${BPF_PROBE_FILENAME}" "${HOME}/.sysdig/${BPF_PROBE_NAME}.o"
		exit $?
	else
		echo "* Failure to find a BPF probe"
		exit 1
	fi
}

ARCH=$(uname -m)
KERNEL_RELEASE=$(uname -r)
SCRIPT_NAME=$(basename "${0}")
SYSDIG_PROBE_URL=${SYSDIG_PROBE_URL:-https://s3.amazonaws.com/download.draios.com}

if [ -z "${SYSDIG_REPOSITORY}" ]; then
	SYSDIG_REPOSITORY="stable"
fi

if [ "${SCRIPT_NAME}" = "sysdig-probe-loader" ]; then
        if [ -z "$SYSDIG_VERSION" ]; then
	    SYSDIG_VERSION=$(sysdig --version | cut -d' ' -f3)
	fi
	PROBE_NAME="sysdig-probe"
	BPF_PROBE_NAME="sysdig-probe-bpf"
	PACKAGE_NAME="sysdig"
elif [ "${SCRIPT_NAME}" = "sysdigcloud-probe-loader" ]; then
	EXEPATH=$(dirname "$(readlink -f "${0}")")
        if [ -z "$SYSDIG_VERSION" ]; then
	    SYSDIG_VERSION=$("${EXEPATH}"/dragent --version)
	fi
	PROBE_NAME="sysdigcloud-probe"
	BPF_PROBE_NAME="sysdigcloud-probe-bpf"
	PACKAGE_NAME="draios-agent"
elif [ "${SCRIPT_NAME}" = "falco-probe-loader" ]; then
        if [ -z "$SYSDIG_VERSION" ]; then
	    SYSDIG_VERSION=$(falco --version | cut -d' ' -f3)
	fi
	PROBE_NAME="falco-probe"
	BPF_PROBE_NAME="falco-probe-bpf"
	PACKAGE_NAME="falco"
else
	echo "This script must be called as sysdig-probe-loader, sysdigcloud-probe-loader, or falco-probe-loader"
	exit 1
fi

if [ "$(id -u)" != 0 ]; then
	echo "Installer must be run as root (or with sudo)."
	exit 1
fi

if ! hash curl > /dev/null 2>&1; then
	echo "This program requires curl"
	exit 1
fi

if [ -v SYSDIG_BPF_PROBE ] || [ "${1}" = "bpf" ]; then
	load_bpf_probe
else
	load_kernel_probe
fi
