#!/bin/sh
#
#	$Id: arptables-noarp-addr,v 1.2 2005/12/27 08:43:11 horms Exp $
#
#	Copyright (C) 2005 Horms <horms@verge.net.au>
#
#	This script manages the arptables entries to prevent NOARP_ADDR
#	from being advertised via ARP. Such an address may
#	be present on the loopback interface as part of an LVS 
#	(www.linuxvirtualserver.org) setup. 
#
#	This is done by dropping incoming arp packets for NOARP_ADDR.
#	And mangling outgoing arp packets from NOARP_ADDR.
#	The later is done for each ARPing interface, and the
#	address chosen as the replacement is the first address on
#	the interface on the same subnet as NOARP_ADDR, and if no
#	such address exists, first address on the interface.
#
#	The arptables rules effected are printed to stdout for referance
#
#	usage: arptables-noarp-addr NOARP_ADDR|NOARP_PREFIX... {start|stop|restart|try-restart|force-reload|status}
#             NOARP_ADDR:   IP address     (e.g. 10.0.0.1) 
#             NOARP_PREFIX: Network prefix (e.g. 10.0.0.1/23)
#
#	BUGS
#	  There seems to be a problem with arptables on Debian at least,
#	  such that removing a rule that includes the mangle target does
#	  not check the --mangle-* argumetns, and will delete non-matching
#	  entries. 
#
unset LANG
LC_ALL=C
export LC_ALL


IPCMD="/sbin/ip"
ARPTABLESCMD="/sbin/arptables"

if [ ! -x "$IPCMD" ]; then
	echo "ERROR: $IPCMD missing" >&2
	exit 5
fi

if [ ! -x "$ARPTABLESCMD" ]; then
	echo "ERROR: $ARPTABLESCMD missing" >&2
	exit 5
fi

if $ARPTABLESCMD --version >& /dev/null; then
	# Seems to be artpables
	VARIANT="arptables"
	OUTCHAIN="OUTPUT"
	INCHAIN="INPUT"
else
	# Seems to be artpables_jf
	VARIANT="arptables_jf"
	OUTCHAIN="OUT"
	INCHAIN="IN"
fi


get_arp_iface ()
{
	$IPCMD -o link sh | while read n iface flag tail; do
		iface="${iface%:}"
		if [ "$flag" != "${flag%[<,]LOOPBACK[>,]*}" -o \
				"$flag" != "${flag%[<,]NOARP[>,]*}" -a \
				"$flag" = "${flag%[<,]SLAVE[>,]*}" ]; then
			continue;
		fi
		echo -n "$iface "
	done
}

get_iface_addr ()
{
	local IFACE="$1"
	$IPCMD -o addr sh $IFACE | while read n iface proto addr tail; do
		if [ "$proto" != "inet" ]; then
			continue;
		fi
		echo -n "$addr "
	done
}

ip_to_num ()
{
	local IP="$1"
	local DEC="0"
	while [ 1 ]; do
		A="${IP#*.}"
		OCTET="${IP%.$A}"
		DEC="$(( $DEC*256 + $OCTET))"
		if [ "$A" = "$IP" ]; then
			break;
		fi
		IP="$A"
	done
	echo $DEC
}

# Note, it seems that DASH uses a signed 32 bit entity for
#       maths, and thus the below may give false negatives.
#       This probably isn't an issue given the way this function is used.
same_subnet ()
{
	local VADDR="$1"
	local ADDR="$2"

	IP="${ADDR%/*}"
	if [ "$IP" = "$ADDR" ]; then
		MASKBITS="32"
	else 
		MASKBITS="${ADDR#$IP/}"
	fi

	VIP="${VADDR%/*}"
	if [ "$VIP" = "$VADDR" ]; then
		VMASKBITS="32"
	else 
		VMASKBITS="${VADDR#$VIP/}"
	fi

	D_IP="$(ip_to_num $IP)"
	D_VIP="$(ip_to_num $VIP)"
	M_IP="$(($D_IP >> ( 32 - $MASKBITS )))"
	M_VIP="$(($D_VIP >> ( 32 - $VMASKBITS )))"
	
	if [ "$M_IP" = "$M_VIP" ]; then
		return 0
	fi

	return 1
}

get_best_addr ()
{
	local VADDR="$1"
	local IFACE="$2"

	local ADDR="$(get_iface_addr $2)"

	if [ -z "$ADDR" ]; then
		return 1
	fi

	# Find the first address that is in the same subnet as the VADDR
	for a in $ADDR; do
		if same_subnet $VADDR $a; then
			echo ${a%/*}
			return 0
		fi
	done

	# No address in the same subnet, just return the first address
	a="${ADDR%% *}"
	echo ${a%/*}
}

parse_arptables ()
{

	local IN_INTERFACE="any"
	local OUT_INTERFACE="any"
	local SOURCE_IP="0.0.0.0/0"
	local DESTINATION_IP="0.0.0.0/0"
	local SOURCE_MAC="00/00"
	local DESTINATION_MAC="00/00"
	local H_LENGTH="any"
	local OPCODE="0000/0000"
	local H_TYPE="0000/0000"
	local PROTO_TYPE="0000/0000"

	while [ $# -gt 1 ]; do
		ARG="$1"; shift
		if [ $# = 0 ]; then return 2; fi
		case $ARG in

			-j|--jump)
				JUMP="$1"; shift
				;;

			-s|--source-ip)
				SOURCE_IP="$1"; shift
				if [ SOURCE_IP = "!" ]; then
					if [ $# = 0 ]; then
						return 2
					fi
					SOURCE_IP="$SOURCE_IP $1"; shift
				fi
				;;

			-d|--destination-ip)
				DESTINATION_IP="$1"; shift
				if [ SOURCE_IP = "!" ]; then
					if [ $# = 0 ]; then
						return 2
					fi
					DESTINATION_IP="$DESTINATION_IP $1"
					shift
				fi
				;;

			--source-mac)
				SOURCE_MAC="$1"; shift
				if [ SOURCE_MAC = "!" ]; then
					if [ $# = 0 ]; then
						return 2
					fi
					SOURCE_MAC="$SOURCE_MAC $1"; shift
				fi
				;;

			--destination-mac)
				DESTINATION_MAC="$1"; shift
				if [ DESTINATION_MAC = "!" ]; then
					if [ $# = 0 ]; then
						return 2
					fi
					DESTINATION_MAC="$DESTINATION_MAC $1"
					shift
				fi
				;;

			-i|--in-interface)
				IN_INTERFACE="$1"; shift
				if [ IN_INTERFACE = "!" ]; then
					if [ $# = 0 ]; then
						return 2
					fi
					IN_INTERFACE="$IN_INTERFACE $1"
					shift
				fi
				;;

			-o|--out-interface)
				OUT_INTERFACE="$1"; shift
				if [ OUT_INTERFACE = "!" ]; then
					if [ $# = 0 ]; then
						return 2
					fi
					OUT_INTERFACE="$OUT_INTERFACE $1"
					shift
				fi
				;;

			-l|--h-length)
				H_LENGTH="$1"; shift
				;;

			--opcode)
				OPCODE="$1"; shift
				;;

			--h-type)
				H_TYPE="$1"; shift
				;;

			--proto-type)
				PROTO_TYPE="$1"; shift
				;;

			--mangle-ip-s)
				MANGLE_IP_S="$1"; shift
				;;

			--mangle-ip-d)
				MANGLE_IP_D="$1"; shift
				;;

			--mangle-mac-s)
				MANGLE_MAC_S="$1"; shift
				;;

			--mangle-mac-d)
				MANGLE_MAC_D="$1"; shift
				;;

			--mangle-target)
				MANGLE_TARGET="$1"; shift
				;;

			*)
				echo "unknown arg: $ARG" >&2
				return 2
				;;

		esac
	done

	RULE=""
	if [ ! -z "$JUMP" ]; then
		RULE="$RULE $JUMP"
	fi
	if [ ! -z "$IN_INTERFACE" ]; then
		RULE="$RULE $IN_INTERFACE"
	fi
	if [ ! -z "$OUT_INTERFACE" ]; then
		RULE="$RULE $OUT_INTERFACE"
	fi
	if [ ! -z "$SOURCE_IP" ]; then
		RULE="$RULE $SOURCE_IP"
	fi
	if [ ! -z "$DESTINATION_IP" ]; then
		RULE="$RULE $DESTINATION_IP"
	fi
	if [ ! -z "$SOURCE_MAC" ]; then
		RULE="$RULE $SOURCE_MAC"
	fi
	if [ ! -z "$DESTINATION_MAC" ]; then
		RULE="$RULE $DESTINATION_MAC"
	fi
	if [ ! -z "$H_LENGTH" ]; then
		RULE="$RULE $H_LENGTH"
	fi
	if [ ! -z "$OPCODE" ]; then
		RULE="$RULE $OPCODE"
	fi
	if [ ! -z "$H_TYPE" ]; then
		RULE="$RULE $H_TYPE"
	fi
	if [ ! -z "$PROTO_TYPE" ]; then
		RULE="$RULE $PROTO_TYPE"
	fi
	if [ ! -z "$MANGLE_IP_S" ]; then
		RULE="$RULE --mangle-ip-s $MANGLE_IP_S"
	fi
	if [ ! -z "$MANGLE_IP_D" ]; then
		RULE="$RULE --mangle-ip-d $MANGLE_IP_D"
	fi
	if [ ! -z "$MANGLE_MAC_S" ]; then
		RULE="$RULE --mangle-mac-s $MANGLE_MAC_S"
	fi
	if [ ! -z "$MANGLE_MAC_D" ]; then
		RULE="$RULE --mangle-mac-d $MANGLE_MAC_D"
	fi
	if [ ! -z "$MANGLE_TARGET" ]; then
		RULE="$RULE --mangle-target $MANGLE_TARGET"
	fi

	echo $RULE
}

do_check_arptables ()
{
	local CHAIN="$1"; shift
	RULE1="$(parse_arptables $@)"
	$ARPTABLESCMD -n -L $CHAIN -v |  \
				sed -e '/^[^-]/ d' -e 's/ ,.*//' \
					-e 's/\*/any/g' | \
				while read line; do
		RULE2="$(parse_arptables $line)" || return $?
		if [ "$RULE1" = "$RULE2" ]; then
			return 1
		fi
	done || return $?
}


do_check_arptables_jf ()
{
	local CHAIN="$1"; shift
	RULE1="$(parse_arptables $@)"
	$ARPTABLESCMD -n -L $CHAIN -v |  \
				sed -e 's/  */ /g' -e 's/\*/any/g' | \
				while read bytes packets tail; do
		if [ "$RULE1" = "$tail" ]; then
			return 1
		fi
	done || return $?
}

check_arptables ()
{
	if [ "$VARIANT" = "arptables" ]; then
		do_check_arptables $*
	else
		do_check_arptables_jf $*
	fi
}

add_arptables ()
{
	check_arptables $*
	RC=$?
	if [ $RC = 1 ] ; then
		: Exists, Skip
		return 0
	elif [ $RC = 0 ] ; then
		echo $ARPTABLESCMD -A $*
		$ARPTABLESCMD -A $*
	else
		: Unknown error
		return $RC
	fi
}

del_arptables ()
{
	while [ 1 ]; do
		check_arptables $*
		RC=$?
		if [ $RC = 0 ] ; then
			: Does not exist, Skip
			return 0
		elif [ $RC = 1 ] ; then
			echo $ARPTABLESCMD -D $*
			$ARPTABLESCMD -D $*
		else
			: Unknown error
			return $RC
		fi
	done
}

do_cmd ()
{
	CMD="$1" ; shift
	VADDR="$*"
	IFACE="$(get_arp_iface)"
	for v in $VADDR; do
		$CMD $INCHAIN -j DROP -d $v || return $?
		for i in $IFACE; do
			a="$(get_best_addr $v $i)" || continue
			$CMD $OUTCHAIN -j mangle -o $i -s $v \
				--mangle-ip-s $a || return $?
		done
	done
}

start ()
{
	do_cmd add_arptables "$*"
}

stop ()
{
	do_cmd del_arptables "$*"
}

get_status ()
{
	do_cmd check_arptables "$*"
	RC=$?
	if [ $RC = 1 ]; then
		return 0
	fi
	return 3
}

status ()
{
	get_status "$*"
	RC=$?
	if [ $RC = 0 ]; then
		echo "$0 is running for $*"
	else 
		echo "$0 is stopped for $*"
	fi
	return $RC
}

usage ()
{

	echo "usage: arptables-noarp-addr NOARP_ADDR|NOARP_PREFIX... {start|stop|restart|try-restart|force-reload|status}" >&2
	echo "	NOARP_ADDR:   IP address     (e.g. 10.0.0.1) " >&2
	echo "	NOARP_PREFIX: Network prefix (e.g. 10.0.0.1/23)" >&2
	exit 1
}

{
	if [ $# -lt 2 ]; then
		usage
	fi
	while [  $# -gt 1 ]; do
		VADDR="$VADDR $1"
		shift
	done
	ACTION="$1"

	case $ACTION in
		start)
			start $VADDR
			;;
		stop)
			stop $VADDR
			;;
		restart)
			stop $VADDR
			start $VADDR
			;;
		force-reload|try-restart)
			if get_status $VADDR; then
				stop $VADDR;
				start $VADDR;
			fi
			;;
		status)
			status $VADDR
			;;
		*)
			usage
			;;
	esac

}


