Robustify script

SMP-analyze.sh

-#!/bin/bash -x
+#!/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
@@ -8,14 +10,38 @@ 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}')
-directory=/tmp/remote-analyze/${md5name}
+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
-	ping -c 5 -w 5 $remotehost
+
+	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
@@ -26,16 +52,22 @@ verify_host_live()
 copy_STARS_info()
 {
 	remotehost=$1
+
+	echo "Retrieve STARS info from host $remotehost"
+
 	# Copy the answer back
-	scp -o ConnectTimeout=75 -o BatchMode=yes -P $IDA_PRO_SERVER_PORT $IDA_PRO_SERVER_USER@$remotehost:${directory}/a.ncexe.* .
+	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=10 -o BatchMode=yes -p $IDA_PRO_SERVER_PORT $IDA_PRO_SERVER_USER@$remotehost ls -l ${directory}/a.ncexe.infoannot > tmp.$$
+	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
@@ -53,6 +85,7 @@ 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`
@@ -64,42 +97,62 @@ copy_and_verify_result()
 	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"
-	ssh -o ConnectTimeout=10 -o BatchMode=yes -p $IDA_PRO_SERVER_PORT $IDA_PRO_SERVER_USER@$remotehost mkdir -p ${directory}
+
+	# 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
-	    while [ `ssh -o ConnectTimeout=10 -o BatchMode=yes -p $IDA_PRO_SERVER_PORT $IDA_PRO_SERVER_USER@$remotehost pgrep idal|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
+		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
 
-	scp -o ConnectTimeout=75 -o BatchMode=yes -P $IDA_PRO_SERVER_PORT -q $@ $IDA_PRO_SERVER_USER@$remotehost:$directory
+	# 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
 
-	ssh -o ConnectTimeout=7200 -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
+	# 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
-	return $?
+	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
@@ -107,7 +160,9 @@ if [  $? -eq 0 ]; then
 		echo SERVER HAS ALREADY ANALYZED $md5name, retrieving cached info
 		copy_and_verify_result $remotehost $file.annot
 		exit_code=$?
-	else
+	fi
+
+	if [ ! $exit_code -eq 0 ]; then
 		run_remote_command $IDA_PRO_SERVER_HOST $@
 		exit_code=$?
 	fi
@@ -115,17 +170,26 @@ 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
 
-	run_remote_command $IDA_PRO_SERVER_HOST2 $@
-	exit_code=$?
+	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