#!/usr/bin/env bash
#
# Copyright (C) 2019 The Falco Authors.
#
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Simple script that desperately tries to load the kernel instrumentation by
# looking for it in a bunch of ways. Convenient when running falco inside
# a container or in other weird environments.
#

#
# Returns 1 if $cos_ver > $base_ver, 0 otherwise
#
cos_version_greater()
{
	if [[ $cos_ver == "${base_ver}" ]]; then
		return 0
	fi

	#
	# COS build numbers are in the format x.y.z
	#
	a=$(echo "${cos_ver}" | cut -d. -f1)
	b=$(echo "${cos_ver}" | cut -d. -f2)
	c=$(echo "${cos_ver}" | cut -d. -f3)

	d=$(echo "${base_ver}" | cut -d. -f1)
	e=$(echo "${base_ver}" | cut -d. -f2)
	f=$(echo "${base_ver}" | cut -d. -f3)

	# Test the first component
	if [[ $a -gt $d ]]; then
		return 1
	elif [[ $d -gt $a ]]; then
		return 0
	fi

	# Test the second component
	if [[ $b -gt $e ]]; then
		return 1
	elif [[ $e -gt $b ]]; then
		return 0
	fi

	# Test the third component
	if [[ $c -gt $f ]]; then
		return 1
	elif [[ $f -gt $c ]]; then
		return 0
	fi

	# If we get here, probably malformatted version string?

	return 0
}

get_kernel_config() {
	if [ -f /proc/config.gz ]; then
		echo "Found kernel config at /proc/config.gz"
		KERNEL_CONFIG_PATH=/proc/config.gz
	elif [ -f "/boot/config-${KERNEL_RELEASE}" ]; then
		echo "Found kernel config at /boot/config-${KERNEL_RELEASE}"
		KERNEL_CONFIG_PATH=/boot/config-${KERNEL_RELEASE}
	elif [ -n "${HOST_ROOT}" ] && [ -f "${HOST_ROOT}/boot/config-${KERNEL_RELEASE}" ]; then
		echo "Found kernel config at ${HOST_ROOT}/boot/config-${KERNEL_RELEASE}"
		KERNEL_CONFIG_PATH="${HOST_ROOT}/boot/config-${KERNEL_RELEASE}"
	elif [ -f "/usr/lib/ostree-boot/config-${KERNEL_RELEASE}" ]; then
		echo "Found kernel config at /usr/lib/ostree-boot/config-${KERNEL_RELEASE}"
		KERNEL_CONFIG_PATH="/usr/lib/ostree-boot/config-${KERNEL_RELEASE}"
	elif [ -n "${HOST_ROOT}" ] && [ -f "${HOST_ROOT}/usr/lib/ostree-boot/config-${KERNEL_RELEASE}" ]; then
		echo "Found kernel config at ${HOST_ROOT}/usr/lib/ostree-boot/config-${KERNEL_RELEASE}"
		KERNEL_CONFIG_PATH="${HOST_ROOT}/usr/lib/ostree-boot/config-${KERNEL_RELEASE}"
	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 -> $HOST_ROOT/lib/modules
		echo "Found kernel config at /lib/modules/${KERNEL_RELEASE}/config"
		KERNEL_CONFIG_PATH="/lib/modules/${KERNEL_RELEASE}/config"
	fi

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

	if [[ "${KERNEL_CONFIG_PATH}" == *.gz ]]; then
		HASH=$(zcat "${KERNEL_CONFIG_PATH}" | md5sum - | cut -d' ' -f1)
	else
		HASH=$(md5sum "${KERNEL_CONFIG_PATH}" | cut -d' ' -f1)
	fi
}

get_target_id() {
	if [ -f "${HOST_ROOT}/etc/os-release" ]; then
		# freedesktop.org and systemd
		# shellcheck source=/dev/null
		source "${HOST_ROOT}/etc/os-release"
		OS_ID=$ID
	elif [ -f "${HOST_ROOT}/etc/debian_version" ]; then
		# Older Debian
		# fixme > can this happen on older Ubuntu?
		OS_ID=debian
	elif [ -f "${HOST_ROOT}/etc/centos-release" ]; then
		# Older CentOS
		OS_ID=centos
	else
		>&2 echo "Detected an unsupported target system, please get in touch with the Falco community"
		exit 1
	fi

	case "${OS_ID}" in
	("amzn")
		if [[ $VERSION_ID == "2" ]]; then
			TARGET_ID="amazonlinux2"
		else
			TARGET_ID="amazonlinux"
		fi
		;;
	("ubuntu")
		if [[ $KERNEL_RELEASE == *"aws"* ]]; then
			TARGET_ID="ubuntu-aws"
		else
			TARGET_ID="ubuntu"
		fi
		;;
	(*)
		TARGET_ID=$(echo "${OS_ID}" | tr '[:upper:]' '[:lower:]')
		;;
	esac
}

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

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

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

	echo "* Unloading ${DRIVER_NAME} module, if present"
	rmmod "${DRIVER_NAME}" 2>/dev/null
	WAIT_TIME=0
	KMOD_NAME=$(echo "${DRIVER_NAME}" | tr "-" "_")
	while lsmod | grep "${KMOD_NAME}" > /dev/null 2>&1 && [ $WAIT_TIME -lt "${MAX_RMMOD_WAIT}" ]; do
		if rmmod "${DRIVER_NAME}" 2>/dev/null; then
			echo "* Unloading ${DRIVER_NAME} module succeeded after ${WAIT_TIME}s"
			break
		fi
		((++WAIT_TIME))
		if (( WAIT_TIME % 5 == 0 )); then
			echo "* ${DRIVER_NAME} module still loaded, waited ${WAIT_TIME}s (max wait ${MAX_RMMOD_WAIT}s)"
		fi
		sleep 1
	done

	if lsmod | grep "${KMOD_NAME}" > /dev/null 2>&1; then
		echo "* ${DRIVER_NAME} module 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
		if hash dkms &>/dev/null && dkms install -m "${DRIVER_NAME}" -v "${DRIVER_VERSION}" -k "${KERNEL_RELEASE}" 2>/dev/null; then
			echo "* Trying to load a dkms ${DRIVER_NAME} module, if present"

			if insmod "/var/lib/dkms/${DRIVER_NAME}/${DRIVER_VERSION}/${KERNEL_RELEASE}/${ARCH}/module/${DRIVER_NAME}.ko" > /dev/null 2>&1; then
				echo "${DRIVER_NAME} module found and loaded in dkms"
				exit 0
			elif insmod "/var/lib/dkms/${DRIVER_NAME}/${DRIVER_VERSION}/${KERNEL_RELEASE}/${ARCH}/module/${DRIVER_NAME}.ko.xz" > /dev/null 2>&1; then
				echo "${DRIVER_NAME} module found and loaded in dkms (xz)"
				exit 0
			else
				echo "* Unable to insmod"
			fi
		else
			DKMS_LOG="/var/lib/dkms/${DRIVER_NAME}/${DRIVER_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 ${DRIVER_NAME} driver, if present"

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

	echo "* Trying to find a prebuilt ${DRIVER_NAME} module for kernel ${KERNEL_RELEASE}"

	get_target_id

	local FALCO_KERNEL_MODULE_FILENAME="${DRIVER_NAME}_${TARGET_ID}_${KERNEL_RELEASE}_${KERNEL_VERSION}.ko"

	if [ -f "${HOME}/.falco/${FALCO_KERNEL_MODULE_FILENAME}" ]; then
		echo "Found a prebuilt module at ${HOME}/.falco/${FALCO_KERNEL_MODULE_FILENAME}, loading it"
		insmod "${HOME}/.falco/${FALCO_KERNEL_MODULE_FILENAME}"
		exit $?
	fi

	local URL
	URL=$(echo "${DRIVERS_REPO}/${DRIVER_VERSION}/${FALCO_KERNEL_MODULE_FILENAME}" | sed s/+/%2B/g)

	echo "* Trying to download prebuilt module from ${URL}"
	if curl -L --create-dirs "${FALCO_DRIVER_CURL_OPTIONS}" -o "${HOME}/.falco/${FALCO_KERNEL_MODULE_FILENAME}" "${URL}"; then
		echo "Download succeeded, loading module"
		insmod "${HOME}/.falco/${FALCO_KERNEL_MODULE_FILENAME}"
		exit $?
	else
		>&2 echo "Download failed, consider compiling your own ${DRIVER_NAME} module and loading it or getting in touch with the Falco 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 [ -n "${HOST_ROOT}" ] && [ -f "${HOST_ROOT}/etc/os-release" ]; then
		# shellcheck source=/dev/null
		source "${HOST_ROOT}/etc/os-release"

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

	if [ -n "${HOST_ROOT}" ] && [ -f "${HOST_ROOT}/etc/VERSION" ]; then
		MINIKUBE=1
		MINIKUBE_VERSION="$(cat "${HOST_ROOT}/etc/VERSION")"
	fi

	get_target_id

	local BPF_PROBE_FILENAME="${DRIVER_NAME}_${TARGET_ID}_${KERNEL_RELEASE}_${KERNEL_VERSION}.o"

	if [ ! -f "${HOME}/.falco/${BPF_PROBE_FILENAME}" ]; then

		local BPF_KERNEL_SOURCES_URL=""
		local STRIP_COMPONENTS=1

		customize_kernel_build() {
			if [ -n "${KERNEL_EXTRA_VERSION}" ]; then
			sed -i "s/LOCALVERSION=\"\"/LOCALVERSION=\"${KERNEL_EXTRA_VERSION}\"/" .config
			fi
			make olddefconfig > /dev/null
			make modules_prepare > /dev/null
		}

		if [ -n "${COS}" ]; then
			echo "* COS detected (build ${BUILD_ID}), using cos kernel headers..."

			BPF_KERNEL_SOURCES_URL="https://storage.googleapis.com/cos-tools/${BUILD_ID}/kernel-headers.tgz"
			KERNEL_EXTRA_VERSION="+"
			STRIP_COMPONENTS=0

			customize_kernel_build() {
				pushd usr/src/* > /dev/null || exit

				# Note: this overrides the KERNELDIR set while untarring the tarball
				KERNELDIR=$(pwd)
				export KERNELDIR

				sed -i '/^#define randomized_struct_fields_start	struct {$/d' include/linux/compiler-clang.h
				sed -i '/^#define randomized_struct_fields_end	};$/d' include/linux/compiler-clang.h

				popd > /dev/null || exit

				# Might need to configure our own sources depending on COS version
				cos_ver=${BUILD_ID}
				base_ver=11553.0.0

				cos_version_greater
				greater_ret=$?

				if [[ greater_ret -eq 1 ]]; then
				export KBUILD_EXTRA_CPPFLAGS=-DCOS_73_WORKAROUND
				fi
			}
		fi

		if [ -n "${MINIKUBE}" ]; then
			echo "* Minikube detected (${MINIKUBE_VERSION}), using linux kernel sources for minikube kernel"
			local kernel_version
			kernel_version=$(uname -r)
			local -r kernel_version_major=$(echo "${kernel_version}" | cut -d. -f1)
			local -r kernel_version_minor=$(echo "${kernel_version}" | cut -d. -f2)
			local -r kernel_version_patch=$(echo "${kernel_version}" | cut -d. -f3)

			if [ "${kernel_version_patch}" == "0" ]; then
				kernel_version="${kernel_version_major}.${kernel_version_minor}"
			fi

			BPF_KERNEL_SOURCES_URL="http://mirrors.edge.kernel.org/pub/linux/kernel/v${kernel_version_major}.x/linux-${kernel_version}.tar.gz"
		fi

		if [ -n "${BPF_USE_LOCAL_KERNEL_SOURCES}" ]; then
			local -r kernel_version_major=$(uname -r | cut -d. -f1)
			local -r kernel_version=$(uname -r | cut -d- -f1)
			KERNEL_EXTRA_VERSION="-$(uname -r | cut -d- -f2)"

			echo "* Using downloaded kernel sources for kernel version ${kernel_version}..."

			BPF_KERNEL_SOURCES_URL="http://mirrors.edge.kernel.org/pub/linux/kernel/v${kernel_version_major}.x/linux-${kernel_version}.tar.gz"
		fi

		if [ -n "${BPF_KERNEL_SOURCES_URL}" ]; then
			echo "* Downloading ${BPF_KERNEL_SOURCES_URL}"

			mkdir -p /tmp/kernel
			cd /tmp/kernel || exit
			cd "$(mktemp -d -p /tmp/kernel)" || exit
			if ! curl -L -o kernel-sources.tgz --create-dirs "${FALCO_DRIVER_CURL_OPTIONS}" "${BPF_KERNEL_SOURCES_URL}"; then
				exit 1;
			fi

			echo "* Extracting kernel sources"

			mkdir kernel-sources && tar xf kernel-sources.tgz -C kernel-sources --strip-components "${STRIP_COMPONENTS}"

			cd kernel-sources || exit
			KERNELDIR=$(pwd)
			export KERNELDIR

			if [[ "${KERNEL_CONFIG_PATH}" == *.gz ]]; then
				zcat "${KERNEL_CONFIG_PATH}" > .config
			else
				cat "${KERNEL_CONFIG_PATH}" > .config
			fi

			echo "* Configuring kernel"
			customize_kernel_build
		fi

		echo "* Trying to compile the eBPF probe (${BPF_PROBE_FILENAME})"

		make -C "/usr/src/${DRIVER_NAME}-${DRIVER_VERSION}/bpf" > /dev/null

		mkdir -p "${HOME}/.falco"
		mv "/usr/src/${DRIVER_NAME}-${DRIVER_VERSION}/bpf/probe.o" "${HOME}/.falco/${BPF_PROBE_FILENAME}"

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

	if [ ! -f "${HOME}/.falco/${BPF_PROBE_FILENAME}" ]; then
		local URL
		URL=$(echo "${DRIVERS_REPO}/${DRIVER_VERSION}/${BPF_PROBE_FILENAME}" | sed s/+/%2B/g)

		echo "* Trying to download a prebuilt eBPF probe from ${URL}"

		curl -L --create-dirs "${FALCO_DRIVER_CURL_OPTIONS}" -o "${HOME}/.falco/${BPF_PROBE_FILENAME}" "${URL}"
	fi

	if [ -f "${HOME}/.falco/${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 degraded. **"
			echo "** Please ensure to run on a kernel with CONFIG_BPF_JIT on.     **"
			echo "******************************************************************"
		fi

		echo "* eBPF probe located, it's now possible to start Falco"

		ln -sf "${HOME}/.falco/${BPF_PROBE_FILENAME}" "${HOME}/.falco/${DRIVER_NAME}-bpf.o"
		exit $?
	else
		echo "* Failure to find an eBPF probe"
		exit 1
	fi
}

ARCH=$(uname -m)
KERNEL_RELEASE=$(uname -r)
KERNEL_VERSION=$(uname -v | sed 's/#\([[:digit:]]\+\).*/\1/')
DRIVERS_REPO=${DRIVERS_REPO:-"@DRIVERS_REPO@"}
if [ -n "$DRIVER_INSECURE_DOWNLOAD" ]
then
	FALCO_DRIVER_CURL_OPTIONS=-fsSk
else
	FALCO_DRIVER_CURL_OPTIONS=-fsS
fi

MAX_RMMOD_WAIT=60
if [[ $# -ge 1 ]]; then
	MAX_RMMOD_WAIT=$1
fi

DRIVER_VERSION="@PROBE_VERSION@"
DRIVER_NAME="@PROBE_NAME@"

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 FALCO_BPF_PROBE ] || [ "${1}" = "bpf" ]; then
	load_bpf_probe
else
	load_kernel_module
fi