#/bin/bash
############################################################################################
## Assumptions:
##    - The EC2 Name on your instance for the haproxy is simply "haproxy.domain.com"
##    - Your member servers are named "websocket-001.domain.com .. websocket-999.domain.com
##    - Your using VPCs and yoru haproxy sits in the public subnet with an EIP
##    - Your websocket servers are listening on 8080
##
## Your root crontab entry could look like this (run every 15 min):
##    */15 * * * * /path/to/update_haproxy_config.sh > /dev/null 2>&1
## 
############################################################################################

export EC2_CERT=/path/to/your/ec2/cert.pem
export EC2_PRIVATE_KEY=/path/to/your/ec2/pk.pem
MAXCONN=1000
TMPTAGFILE=/tmp/websocket_tags
TMPINSTANCES=/tmp/websocket_instances
TMPHOSTS=/tmp/hosts
TMPHACONFIG=/tmp/testconfig
HACONFIG=/etc/haproxy/haproxy.cfg
HACONFIGSERVERS=/tmp/haconfig
echo "" > $HACONFIGSERVERS
echo "" > $TMPHOSTS
echo "" > $HACONFIGSERVERS

##################################################################
## Query the AWS EC2 API looking for servers named "socketio"
## You can adjust this if you want to name your member servers
## something different
##################################################################

echo "Getting EC2 Data..."
ec2-describe-tags | grep Name | grep websocket | grep -v haproxy > $TMPTAGFILE
ec2-describe-instances | grep INSTANCE  > $TMPINSTANCES
while read line; do
    NAME1=`echo $line | awk '{print $5}' | tr '.' ' ' | awk '{print $1}'`
    NAME2=`echo $line | awk '{print $5}'`
    INSTANCEID=`echo $line | awk '{print $3}'`
    INTERNALIP=`cat $TMPINSTANCES | grep $INSTANCEID | awk '{print $15}'`
    echo  "$INTERNALIP   $NAME1    $NAME2" >> $TMPHOSTS
    echo  "     server     $NAME1 $NAME2:8080 maxconn $MAXCONN weight 10 check" >> $HACONFIGSERVERS
    echo "" >> $HACONFIGSERVERS
done < $TMPTAGFILE

#############################################################
## this modifies the /etc/hosts file so we can
## use the internal ip (10.x.x.x) for traffic rather then
## then EIP
#############################################################
if [ `cat $TMPHOSTS | wc -l` -gt 0 ]; 
then
  echo "Writing New Host File..."
	echo "127.0.0.1    localhost   haproxy  haproxy.domain.com" > /etc/hosts
	cat $TMPHOSTS | sed '/^$/d' | sort -k 2  >> /etc/hosts
fi

#############################################################
##  This rewrites the haproxy.cfg file placing additional 
##  servers between the two blocks:
##      #SERVERBEGIN
##          ....
##      #SERVEREND
##
##  And then reloads the config
#############################################################
if [ `cat $HACONFIGSERVERS | wc -l` -gt 0 ]; 
then
	echo "Writing HA Proxy Config File..."
	SUBSERVER=`cat $HACONFIGSERVERS | sort -u | sed '/^$/d' | sed -e 's/[\/&]/\\&/g'`
	cat $HACONFIG | sed -e "/#SERVERBEGIN/,/#SERVEREND/c\\\t\#SERVERBEGIN\n\tSERVERPLACEHOLDER\n\t\#SERVEREND" | perl -p -i -e "s/SERVERPLACEHOLDER/${SUBSERVER}/g" > $TMPHACONFIG
	cat $TMPHACONFIG > $HACONFIG

	echo "Reloading HAProxy..."
	/etc/init.d/haproxy reload
fi

#############################################################
##  file cleanup
#############################################################
rm -f $TMPTAGFILE
rm -f $TMPINSTANCES
rm -f $HACONFIGSERVERS
rm -f $TMPHOSTS
rm -f $TMPHACONFIG