#!/bin/bash 

# assumption for CGC: only 1 file to be analyzed

if [ -z "$IDA_PRO_SERVER_HOST" ]; then echo Failed to set IDA_PRO_SERVER_HOST; exit 2; fi
if [ -z "$IDA_PRO_SERVER_USER" ]; then
    IDA_PRO_SERVER_USER=`whoami`
fi
if [ -z "$IDA_PRO_SERVER_PORT" ]; then
    IDA_PRO_SERVER_PORT=22
fi

if [ -z "$MAX_IDA_PROCESSES" ]; then
	MAX_IDA_PROCESSES=25
	echo "Max analysis processes not specified, set to $MAX_IDA_PROCESSES"
fi

if [ -z "$1" ]; then
	echo Failed to specify binary to be analyzed
	exit 3
fi

file=$1

md5name=$(md5sum $file | awk '{print $1}')
if [ -z "$md5name" ]; then 
	echo Failed to obtain md5 hash for file to be analyzed
	exit 4
fi

DIRECTORY=/tmp/remote-analyze/${md5name}

# in seconds
COPY_TIMEOUT=75
SIMPLE_TIMEOUT=10
ANALYZE_TIMEOUT=7200

verify_host_live()
{
	remotehost=$1

	echo "Verify host $remotehost is live"

	ssh -o ConnectTimeout=$SIMPLE_TIMEOUT -o BatchMode=yes -p $IDA_PRO_SERVER_PORT $IDA_PRO_SERVER_USER@$remotehost ls /tmp >/dev/null
	if [ ! $? -eq 0 ]; then
		return 1
	fi
	
	return 0
}

copy_STARS_info()
{
	remotehost=$1

	echo "Retrieve STARS info from host $remotehost"

	# Copy the answer back
	scp -o ConnectTimeout=$COPY_TIMEOUT -o BatchMode=yes -P $IDA_PRO_SERVER_PORT $IDA_PRO_SERVER_USER@$remotehost:${DIRECTORY}/a.ncexe.* .
	return $?
}

server_has_cached_info()
{
	remotehost=$1

	echo "Check whether STARS results already present from host $remotehost"

	# Copy the answer back
	ssh -o ConnectTimeout=$SIMPLE_TIMEOUT -o BatchMode=yes -p $IDA_PRO_SERVER_PORT $IDA_PRO_SERVER_USER@$remotehost ls -l ${DIRECTORY}/a.ncexe.infoannot > tmp.$$
	if [ -f tmp.$$ ]; then
		grep a.ncexe.infoannot tmp.$$
		if [ ! $? -eq 0 ]; then
			rm tmp.$$
			return 1
		else
			return 0
		fi
	fi
	return 1
}

copy_and_verify_result()
{
	remotehost=$1
	annotfile=$2

	echo "Copy and sanity check STARS results from host $remotehost"
	copy_STARS_info $remotehost

	lines=`cat $annotfile | wc -l`
	if [ $lines -lt 10 ]; then
		echo "Failed to produce a valid annotations file for $annotfile."
		return 1 
	fi

	return 0
}

cleanup_remote_host()
{
	remotehost=$1
	echo "Cleanup host $remotehost"
	ssh -o ConnectTimeout=$SIMPLE_TIMEOUT -o BatchMode=yes -p $IDA_PRO_SERVER_PORT $remotehost rm -f ${DIRECTORY}/a.i64
}

run_remote_command()
{
	remotehost=$1
	shift
	echo "Remote analyze on host $remotehost"

	# Create remote DIRECTORY
	ssh -o ConnectTimeout=$SIMPLE_TIMEOUT -o BatchMode=yes -p $IDA_PRO_SERVER_PORT $IDA_PRO_SERVER_USER@$remotehost mkdir -p ${DIRECTORY}

	# Check to see if the Ida Pro Server is too busy now and wait if necc.
	if [ -n "$MAX_IDA_PROCESSES" ]; then
		if [ `ssh -o ConnectTimeout=$SIMPLE_TIMEOUT -o BatchMode=yes -p $IDA_PRO_SERVER_PORT $IDA_PRO_SERVER_USER@$remotehost pgrep idal64|wc -l` -ge "$MAX_IDA_PROCESSES" ]; then
			echo "Remote server $remotehost too busy (max allowed: $MAX_IDA_PROCESSES) -- skip it"
			return 1
		fi
#	    while [ `ssh -o ConnectTimeout=$SIMPLE_TIMEOUT -o BatchMode=yes -p $IDA_PRO_SERVER_PORT $IDA_PRO_SERVER_USER@$remotehost pgrep idal64|wc -l` -ge "$MAX_IDA_PROCESSES" ]; do
#	        random=`od -An -N2 -tu2 /dev/urandom`
#	        # Wait 10-30 seconds
#		seconds=`expr $random % 20 + 10`
#		echo Waiting $seconds seconds for an IDA process to exit...
#		sleep $seconds
#	    done
	fi

	# Copy files to be processed to remote DIRECTORY
	scp -o ConnectTimeout=$COPY_TIMEOUT -o BatchMode=yes -P $IDA_PRO_SERVER_PORT -q $@ $IDA_PRO_SERVER_USER@$remotehost:$DIRECTORY
	if [ ! $? -eq 0 ]; then
		return 1
	fi

	# Give it max 2 hours for analysis
	ssh -o ConnectTimeout=$ANALYZE_TIMEOUT -o BatchMode=yes -p $IDA_PRO_SERVER_PORT $IDA_PRO_SERVER_USER@$remotehost "cd ~/techx_umbrella/peasoup; source set_env_vars; export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:.; cd $DIRECTORY; screen -D -L -ln -m -a -T xterm sh -x "'$SMPSA_HOME'"/SMP-analyze.sh $@" 2>&1
	if [ ! $? -eq 0 ]; then
		return 1
	fi

	copy_and_verify_result $remotehost $file.annot
	rc=$?

	# cleanup_remote_host remote directory (a.i64)
	cleanup_remote_host $remotehost
	return $rc
}

# FIXME: need to handle multiple files to be analyzed?

exit_code=1

# try host 1
verify_host_live $IDA_PRO_SERVER_HOST
if [  $? -eq 0 ]; then
	server_has_cached_info $IDA_PRO_SERVER_HOST
	if [ $? -eq 0 ]; then
		echo SERVER HAS ALREADY ANALYZED $md5name, retrieving cached info
		copy_and_verify_result $remotehost $file.annot
		exit_code=$?
	fi

	if [ ! $exit_code -eq 0 ]; then
		run_remote_command $IDA_PRO_SERVER_HOST $@
		exit_code=$?
	fi
else
	echo "Host $IDA_PRO_SERVER_HOST is not responding"
fi

# if needed, try host 2
if [ ! $exit_code -eq 0 ]; then
	echo "Error detected on primary $IDA_PRO_SERVER_HOST, failing over to $IDA_PRO_SERVER_HOST2"
	verify_host_live $IDA_PRO_SERVER_HOST2
	if [ ! $? -eq 0 ]; then
		echo "Backup host $IDA_PRO_SERVER_HOST2 is not responding"
		exit 1
	fi

	server_has_cached_info $IDA_PRO_SERVER_HOST2
	if [ $? -eq 0 ]; then
		echo SERVER HAS ALREADY ANALYZED $md5name, retrieving cached info
		copy_and_verify_result $IDA_PRO_SERVER_HOST2 $file.annot
		exit_code=$?
	fi

	if [ ! $exit_code -eq 0 ]; then
		run_remote_command $IDA_PRO_SERVER_HOST2 $@
		exit_code=$?
	fi
fi

exit $exit_code