# -*- shell-script -*-
# Ver.0.3.0
 
function cidr2oct () {
    # Convert CIDR netmask to x.x.x.x format.
    local mask bit octs i
    mask=$1

    if grep -q '\.' <<<$mask; then
	echo $mask
	return
    fi

    for ((i=$mask; $i>0; i--)); do
	bit="${bit}1"
    done
    i=$((32 - $mask))
    for ((i=$i; $i>0; i--)); do
	bit="${bit}0"
    done

    octs=$(echo 'ibase=2;obase=A;'$(cut -c 1-8 <<<$bit) |bc)
    octs=${octs}.$(echo 'ibase=2;obase=A;'$(cut -c 9-16 <<<$bit) |bc)
    octs=${octs}.$(echo 'ibase=2;obase=A;'$(cut -c 17-24 <<<$bit) |bc)
    octs=${octs}.$(echo 'ibase=2;obase=A;'$(cut -c 25-32 <<<$bit) |bc)

    echo $octs
}

function oct2cidr () {
    # Convert decimal netmask to CIDR.
    local mask bit cidr i
    mask=$1

    if grep -qv '\.' <<<$mask; then
	echo $mask
	return
    fi

    for i in 1 2 3 4; do
	bit=${bit}$(printf "%08d" \
	    $(echo 'ibase=10;obase=2;'$(cut -d '.' -f $i <<<$mask) |bc))
    done
    cidr=$(echo -n ${bit%%0*} |wc -m)

    echo $cidr
}

function set_virnet_vars () {
    ## Auto-detect VirtNet specifications.
    local i j addr net fwdmode fwddev br
    #local mask

    virbr_active=($(virnetinfo.py list))

    i=0; j=0
    for net in ${virbr_active[@]}; do
	fwdmode=$(virnetinfo.py fwdmode $net)
	br=$(virnetinfo.py brname $net)
	if [ "$fwdmode" = "Unknown" ]; then
	    virbr_no_nat[i]=$br
	    : $((i+=1))
	    continue
	fi

	virbrnet_nat[j]=$(ip route show dev $br proto kernel |cut -d' ' -f1)

	## Another way to learn virbrnet_nat, but this needs ipcalc.
	#read addr mask < <(virnetinfo.py addr $net)
	#mask=$(oct2cidr $mask |cut -d'=' -f2)
	#virbrnet_nat[j]=$(ipcalc -s -n ${addr}/${mask} |cut -d'=' -f2)/$mask

	addr=""
	fwddev=$(virnetinfo.py fwddev $net)
	if [ "$fwddev" = "Unknown" ]; then
	    fwddev=$(ip route show 0/0 |cut -d' ' -f5)
	fi
	addr=$(echo $(ip addr show dev $fwddev |grep inet) | \
	    cut -d' ' -f2)
	peth_addr[j]=${addr%%/*}
	: $((j+=1))
    done
}

function disable_ipfwd () {
    ## Disable kernel IP forwarding if there is no NAT bridges.
    if [ ${#virbr_no_nat[@]} -gt 0 -a ${#virbrnet_nat[@]} -lt 1 ]; then
	echo 0 >/proc/sys/net/ipv4/ip_forward
    fi
}

function is_validline () {
    ## FALSE if it seems a comment or blank line, TRUE otherwise.
    grep -Eq '^[[:space:]]*(#|$)' <<<$1 && return 1
    return 0
}

function set_route () {
    ## Add or replace routing tables.
    ## Arg: a file describing routes and/or rules.
    local line
    [ -r "$1" ] || return

    while read line; do
	if is_validline $line; then
	    if grep -Eq '^([0-9]|default)' <<<$line; then
		ip route replace $line
		continue
	    fi

	    ip rule add $line
	fi
    done < "$1"
}

function optimize_ipt () {
    ## Iptables optimization.
    local i net rule omatch rulenum iface

    # Replace MASQUERADE with SNAT rule.
    if [ ${#virbrnet_nat[@]} -gt 0 ]; then
	for i in $(seq 0 $((${#peth_addr[@]} - 1))); do
	    echo ${peth_addr[$i]} |grep -Eq '^([0-9]{1,3}\.){3}[0-9]{1,3}$'
	    [ $? -eq 0 ] || continue

	    net=${virbrnet_nat[$i]}
	    rule=$(iptables-save |grep MASQUERADE |
		grep -F -e "-s ${net%%/*}/$(cidr2oct ${net##*/})" |head -n 1)
	    [ -z "$rule" ] && continue

	    omatch=$(echo "$rule" |sed 's/.*\(-o \w\+\).*/\1/' 2>/dev/null)
	    rulenum=$(iptables -t nat -L -n --line-numbers |
		awk -v n=${net%%/*}/$(oct2cidr ${net##*/}) '{
		  if($2 == "MASQUERADE" && $5 == n){ print $1; exit; }
	    }')

	    iptables -t nat -R POSTROUTING $rulenum \
	      -s $net -d ! $net $omatch -j SNAT --to-source ${peth_addr[$i]}
	done
    fi

    # Delete unnecessary DNS and DHCP rules.
    for iface in ${virbr_no_nat[@]}; do
	iptables -D INPUT \
	    -i $iface -p udp -m udp --dport 53 -j ACCEPT &>/dev/null
	iptables -D INPUT \
	    -i $iface -p tcp -m tcp --dport 53 -j ACCEPT &>/dev/null
	iptables -D INPUT \
	    -i $iface -p udp -m udp --dport 67 -j ACCEPT &>/dev/null
	iptables -D INPUT \
	    -i $iface -p tcp -m tcp --dport 67 -j ACCEPT &>/dev/null
    done
}
