#!/bin/sh
#
# nShaper v0.3
#
# Author: Nikita Popov <nikus@rambler.ru>
# Release date: 22-09-2009
# License: GPL
#
# This script needs a fresh IMQ to be installed
# (Oleg's firmware 1.9.2.7-d-r527 and later, check http://wl500g.info)
# Latest beta patch for other firmwares can be obtained here
# (tested by me for a few weeks, works fine):
# 	http://wiki.nix.hu/cgi-bin/twiki/view/IMQ/ImqDevelImqBehaviour
#
# This script configures Linux routing service using HTB classes and
# [E]SFQ queues. It accepts diferent local WAN zones (e.g. provider's net and
# peerings) with their download rates, so local traffic won't be able to
# interfere with Internet ones. Also the script splits Internet traffic to 5
# different queues (both on input and output) by priority, so realtime apps
# and web serfing won't be disturbed by download managers, even running
# on router.

# This script is written as service. Default path is /opt/etc/init.d/nshaper
# Text files with IPs for local zones are searched
# as /opt/etc/nshaper/ip_[zone name].lst
# Syntax is simple: each ip with subnet mask in form of a.b.c.d/mask
# should be written in separate line, empty lines and lines beginning with '#'
# are ignored. Subnet mask could be written as prefix - /24 or
# as full record - /255.255.255.0

# If your provider switches download rate for some zones, as mine
# (e.g. from day to night) then use crontab. (see comments below)

# If you do not have a USB stick connected to router, it is possible to
# store this script and zone files into flashfs memory space.
# Check http://wl500g.info for more info.
#

#                      * * * GLOBAL CONSTANTS * * *
# WAN_IF (ppp0) Internet connection interface name
WAN_IF=ppp0
# LAN _IF (br0) This interface includes ethernet and wifi
LAN_IF=br0

#                       * * * RATE SETTINGS * * *
# All rates below are in kilobits (kbit)
# WAN port real rate (dsl or ppp download/upload)
WAN_DN_RATE=8192
WAN_UP_RATE=1560

# Zone order represents their priority, decreasing from the first to the last
# "inet" zone means everything not matched with other zones
WAN_ZONES="inet cn z2 z0"

# !!! All rates need to be measured first !!!
# Set rate to value 5-10% less than measured, if real rate is less than rate
# in list below, the shaper will not work properly!
# I've got double speed for inet and cn zones at night between 1:00-8:00
if test `date +%H` -ge 1 -a `date +%H` -lt 8; then
  WAN_ZONES_DN_RATE="1900 1900 1900 6200"
  WAN_ZONES_UP_RATE="1500 1500 1500 1500"
else
  WAN_ZONES_DN_RATE=" 950  950 1900 6200"
  WAN_ZONES_UP_RATE="1500 1500 1500 1500"
fi

# Shaper creates 5 bands for Internet traffic (see schemes below)
# Guarantee rate for each band is, in %-s:
RATES="10 20 50 10 10"


#######################################################################
#                 S H A P E R         S C H E M A T I C               #
#######################################################################
#                                                                     #
#                             prerouting chain                        #
#                            +-----------+-----+                      #
#             +------+ IMQ0  |  tc LAN   | sfq |  IMQ0                #
#         +-->| DNAT |------>|  ip:port  | ... |----------->+         #
#         |   +------+  in   |  filters  | sfq |  out       |         #
#    WAN  |                  +-----------+-----+            |         #
#   <====>|                        FILTERS                  |         #
#    ppp0 |                  +-----+-----------+            |         #
#     or  |   +------+  IMQ1 | sfq |tc LAN/WAN |  IMQ1      |         #
#   vlan1 +<--| SNAT |<------| ... |  ip:port  |<----------+|         #
#             +------+  out  | sfq |  filters  |  in       ||         #
#                            +-----+-----------+           ||         #
#                            postrouting chain             ||         #
#                                                          ||         #
#         +--------------+     +--------------+------------||+        #
#         |              |---->| output chain |   forward    |        #
#         |  Local apps  |     +==============+              |        #
#         |              |<----| input  chain |    chain     |        #
#         +--------------+     +--------------+------------||+        #
#                                       routing core       ||         #
#                                                     LAN  ||         #
#                                                    <====>++         #
#                                                      br0            #
#                                                                     #
#######################################################################


# Initial conditions:
# 1. There are a number of local zones on WAN interface
#    (provider's network, peerings with another providers, and so on).
# 2. Rates of some zones can be varied during the day (e.g. because of plan).
# 3. Internet traffic should take priority of all other traffic zones.
# 4. Internet traffic should be divided between users and local apps
#    (such as P2P) by our rules. Interactive traffic should be forced.
# 5. Local interactive services (ssh, web) should be forced, and downloaders
#    should take only free bandwidth.
# 6. Time-critical apps should always work perfectly. :)


# clear previous setup
clear_cfg() {
  iptables -t mangle -D PREROUTING -i $WAN_IF -j IMQ --todev 0 > /dev/null 2>&1
  tc qdisc del dev imq0 root > /dev/null 2>&1

  iptables -t mangle -D POSTROUTING -o $WAN_IF -j IMQ --todev 1 > /dev/null 2>&1
  tc qdisc del dev imq1 root > /dev/null 2>&1
}

# clear setup and unload modules
stop() {
  clear_cfg
  ip link set dev imq0 down > /dev/null 2>&1
  ip link set dev imq1 down > /dev/null 2>&1
  rmmod imq > /dev/null 2>&1
  rmmod ipt_IMQ > /dev/null 2>&1
  rmmod sch_esfq > /dev/null 2>&1
}

# start nShaper
start() {

# IP addresses on interfaces
WAN_IP=`nvram get wan_ipaddr_t 2>/dev/null`
LAN_IP=`nvram get lan_ipaddr_t 2>/dev/null`
LAN_MASK=`nvram get lan_netmask_t 2>/dev/null`
MTU=`nvram get wan_pppoe_mtu 2>/dev/null`

# Alternate method of getting interface data
if test -z "$WAN_IP" -o -z "$MTU" -o -z "$LAN_IP" -o -z "$LAN_MASK"; then
  tmp=`ifconfig $WAN_IF`
  tmp=${tmp#*inet addr:}
  WAN_IP=${tmp%% *}

  tmp=`ifconfig $LAN_IF`
  tmp=${tmp#*inet addr:}
  LAN_IP=${tmp%% *}

  tmp=`ifconfig $LAN_IF`
  tmp=${tmp#*Mask:}
  LAN_MASK=${tmp%%$'\n'*}

  tmp=`ip link show $WAN_IF`
  tmp=${tmp#*mtu }
  MTU=${tmp%% *}

  #echo "wan ip: $WAN_IP mtu: $MTU"
  #echo "lan ip: $LAN_IP mask: $LAN_MASK"
fi


# comment modules checking/loading (and unloading in "stop" function)
# if your kernel compiled with imq/esfq code

# get list of modules loaded
MODULES=`lsmod`

# check for IMQ module
if test -n "${MODULES%%*imq*}"; then
  insmod ipt_IMQ > /dev/null 2>&1
  insmod imq behaviour=ab numdevs=2 > /dev/null 2>&1
  if test $? -ne 0; then
    echo "Can't load imq module. Update your firmware first."
    exit 1
  fi
  ip link set dev imq0 up > /dev/null 2>&1
  ip link set dev imq1 up > /dev/null 2>&1
  ip link set dev imq0 mtu $MTU > /dev/null 2>&1
  ip link set dev imq1 mtu $MTU > /dev/null 2>&1
fi

# check for esfq module
QUEUE="esfq quantum $MTU limit 128 depth 128 hash dst"
if test -n "${MODULES%%*sch_esfq*}"; then
  insmod /opt/lib/modules/2.4.37.5/sch_esfq.o > /dev/null 2>&1
  if test $? -ne 0; then
    echo "Can't find ESFQ module. Using SFQ queues..."
    QUEUE="sfq quantum $MTU perturb 10"
  fi
fi


# Shell array implememtation
# echoes element from array
# Usage: `array_pop index array_el[0] array_el[1] ... array_el[n-1]`
array_pop() {
  num=$(($1+1))
  shift $num
  echo $1
}

#############################################################################
#                         Incoming traffic shaper                           #
#############################################################################


#   WAN/MAN
#  =========>+--- Internet (highest priority)
#   traffic  |       |
#            |       +-- Time-critical queue
#            |       |   (TCP ACKs, DNS, ICMP, TOS:Low delay)
#            |       |
#            |       +-- High-priority queue (VOIP, MMS, games, etc)
#            |       |
#            |       +-- Middle-priority queue (HTTP(S), AIM)
#            |       |
#            |       +-- Low-priority queue (all others - e.g. FTP, MAIL)
#            |       |
#            |       +-- Lazy queue (P2P)
#            |
#            +---  Zone 0 (lower priority)
#            |
#            +---   ...
#            |
#            +---  Zone N (lowest priority)

echo "Setup downlink shaper "
# create root class
tc qdisc add dev imq0 root handle 1:0 htb
tc class add dev imq0 parent 1:0 classid 1:1 htb \
    rate ${WAN_DN_RATE}kbit burst 25k

# setup zones
bw=$WAN_DN_RATE
CLSID_INET=""
i=0
for rate in $WAN_ZONES_DN_RATE; do
  if test $((bw - rate)) -lt 0; then
    # not enough bandwidth - limiting rate
    tc class add dev imq0 parent 1:1 classid 1:2${i} htb \
      rate ${bw}kbit ceil ${rate}kbit prio ${i}0
    bw=1
  else
    # add zone class
    tc class add dev imq0 parent 1:1 classid 1:2${i} htb \
      rate ${rate}kbit prio ${i}0
    bw=$((bw - rate))
  fi

  if test "`array_pop $i $WAN_ZONES`" = "inet"; then
    # found "inet" entry!
    # create 5 bands for Internet traffic
    for j in `seq 0 4`; do
      tc class add dev imq0 parent 1:2${i} classid 1:2${i}${j} htb \
        rate $((`array_pop $j $RATES`*rate/100))kbit \
        ceil $((9*rate/10))kbit prio ${j}
      tc qdisc add dev imq0 parent 1:2${i}${j} handle 2${i}${j} $QUEUE
    done
    CLSID_INET="2${i}"
  else
    # add queue for each zone (except "inet")
    tc qdisc add dev imq0 parent 1:2${i} handle 2${i} $QUEUE
  fi

  i=$((i+1))
done

if test -z "$CLSID_INET"; then
  echo "Zone \"inet\" was not found in WAN_ZONES list! Exiting..."
  stop
  exit 1
fi


#############################################################################
#                         Outgoing traffic shaper                           #
#############################################################################

#                                                         WAN/MAN
#   Time-critical queue -------------------------------+=========>
#   (TCP ACKs, DNS, ICMP, TOS:Low delay)               |  traffic
#                                                      |
#   High-priority queue (VOIP, MMS, games, etc) -------+
#                                                      |
#   Middle-priority queue (HTTP(S), AIM) --------------+
#                                                      |
#   Low-priority queue (all others - e.g. FTP, MAIL)---+
#                                                      |
#   Lazy queue (P2P)-----------------------------------+
#

echo "Setup uplink shaper "

# create root class
tc qdisc add dev imq1 root handle 2:0 htb
tc class add dev imq1 parent 2:0 classid 2:1 htb \
    rate ${WAN_UP_RATE}kbit mtu $MTU

# setup zones
bw=$WAN_UP_RATE
i=0
for rate in $WAN_ZONES_UP_RATE; do
  if test $((bw - rate)) -lt 0; then
    # not enough bandwidth - limiting rate
    tc class add dev imq1 parent 2:1 classid 2:2${i} htb \
      rate ${bw}kbit ceil ${rate}kbit prio ${i}0
    bw=1
  else
    # add zone class
    tc class add dev imq1 parent 2:1 classid 2:2${i} htb \
      rate ${rate}kbit prio ${i}0
    bw=$((bw - rate))
  fi

  if test "`array_pop $i $WAN_ZONES`" = "inet"; then
    # found "inet" entry!
    # create 5 bands for Internet traffic
    for j in `seq 0 4`; do
      tc class add dev imq1 parent 2:2${i} classid 2:2${i}${j} htb \
        rate $((`array_pop $j $RATES`*rate/100))kbit \
        ceil $((9*rate/10))kbit prio ${j}
      tc qdisc add dev imq1 parent 2:2${i}${j} handle 2${i}${j} $QUEUE
    done
  else
    # add queue for each zone (except "inet")
    tc qdisc add dev imq1 parent 2:2${i} handle 2${i} $QUEUE
  fi

  i=$((i+1))
done


#############################################################################
#                              Filters stuff                                #
#############################################################################

# function convertes mask to prefix:	 a.b.c.d/e.f.g.h => a.b.c.d/p
#					 a.b.c.d/p       => a.b.c.d/p
#					 a.b.c.d	 => a.b.c.d
toprefixmask() {
  tmp=$1
  # search for ip/mask delimiter 
  if ! test "$tmp" = "${tmp#*/}"; then
    mask=${tmp#*/}
    # check for dots in mask
    if test "$mask" != "${mask%%.*}"; then
      pfx=`ipcalc -p "${tmp%/*}" "${tmp#*/}"`
      pfx=${pfx##*=}
      tmp="${tmp%/*}/$pfx"
    fi
  fi
  echo "$tmp"
  return 0
}

# function "setrule" makes rules setup much easier
# Usage: setrule [all] [proto tcp|udp|icmp] [short] [flag ack|syn|rst|...] \
#                [afc 1-4] [afd 1-3] \
#                [lan|wan|both] [ip a.b.c.d/m] [port n] [prio u16] \
#        queue 0-4
#
# list of parameters:
#
# all (keyword) 	 - match everything (no filtering, move all to queue)
# proto			 - target proto (icmp/tcp/udp)
# short (keyword)	 - match only short packets (<64b length, =20b header)
# flag ack|syn|rst|...	 - match tcp control flag. may be used multiple times
# afc 1-4		 - match assured forward class (higher => lower latency)
# afd 1-3		 - match assured forward drop prec (lower => lower drop)
# lan/wan/both (keyword) - assignes a place of matched target. default is {both}
# ip ip_addr/mask	 - target ip address
# port port		 - target port
# prio			 - sequence number for filter. auto incremented
#                          if omitted. default is {10}, 1 is reserved!
# queue 0-4              - assignes queue for the rule, lower is better

setrule() {
  rule="$*"
  if test -z "$str_p"; then
    str_p="10"
  else
    str_p=$((str_p+1))
  fi
  str_q=""
  str_m=""
  str_dir="both"
  flag_both=1
  while test -n "$1"; do
    case $1 in
      prio)	str_p="$2"
      		shift 2;;

      lan|wan)	str_dir="$1"
      		flag_both=0
      		shift 1;;

      both)	str_dir="$1"
      		shift 1;;

      all)	str_m=" match u32 0 0 "
      		shift 1;;

      ip)	if test "$str_dir" = "lan"; then
      		   str_m="$str_m match ip dst `toprefixmask $2` "
      		else
      		   str_m="$str_m match ip src `toprefixmask $2` "
      		fi
      		shift 2;;

      port)	if test "$str_dir" = "lan"; then
      		  str_m="$str_m match ip dport $2 0xffff "
      		else
      		  str_m="$str_m match ip sport $2 0xffff "
      		fi
      		shift 2;;

      short)	str_m="$str_m match u8 0x05 0x0f at 0 \
      		              match u16 0x0000 0xffc0 at 2 "
      		shift 1;;

      proto)	tmp=$2
      		if test -z "${tmp%%*[^0-9]*}"; then
      		  case $tmp in
      		    icmp) tmp=1;;
      		    tcp)  tmp=6;;
      		    udp)  tmp=17;;
      		    *)    echo "Unknown  protocol \"$2\" at rule \"$rule\""
      		          return;;
      		  esac
      		fi
      		str_m="$str_m match ip protocol $tmp 0xff "
      		shift 2;;

      flag)	if test -n "${str_m%%*protocol 6*}"; then
      		  str_m="$str_m match ip protocol 6 0xff "
      		fi
      		case $2 in
      		  fin)  str_m="$str_m match u8 0x01 0x01 at nexthdr+13 ";;
      		  syn)  str_m="$str_m match u8 0x02 0x02 at nexthdr+13 ";;
      		  rst)  str_m="$str_m match u8 0x04 0x04 at nexthdr+13 ";;
      		  psh)  str_m="$str_m match u8 0x08 0x08 at nexthdr+13 ";;
      		  ack)  str_m="$str_m match u8 0x10 0x10 at nexthdr+13 ";;
      		  urg)  str_m="$str_m match u8 0x20 0x20 at nexthdr+13 ";;
      		  *)	echo "Bad flag parameter \"$2\" at rule \"$rule\""
      			return;;
      		esac
      		shift 2;;

      afc)	if test $2 -ge 1 -a $2 -le 4; then
      		  str_m="$str_m match u8 "$(($2 * 32))" 0xe0 at 1 "; 
      		else
      		  echo "Wrong afc parameter \"$2\" at rule \"$rule\""
      		  echo "    Assured forwarding class range is 1..4"
      		fi
      		shift 2;;

      afd)	if test $2 -ge 1 -a $2 -le 3; then
      		  str_m="$str_m match u8 "$(($2 * 4))" 0x16 at 1 "; 
      		else
      		  echo "Wrong afd parameter \"$2\" at rule \"$rule\""
      		  echo "    Assured forwarding drop precedence range is 1..3"
      		fi
      		shift 2;;

      queue)	if test $2 -ge 0 -a $2 -le 4; then
      		  str_q="$2"
      		else
      		  echo "Wrong queue \"$2\" at rule \"$rule\""
      		  return
      		fi
      		shift 2;;

      *)	echo "Syntax error: \"$1\" at rule \"$rule\""
      		return;;
    esac
  done

  if test -z "$str_m" -o -z "$str_q"; then
    echo "Missing parameters at rule \"$rule\""
    return
  fi

  # str_m - match string for imq0
  # make reverse string for imq1, where src<=>dst, sport<=>dport
  tmp=${str_m}
  # backup src/sport
  tmp=${tmp//src/sr_c}
  tmp=${tmp//sport/spor_t}
  # fulltext replace dst=>src, dport=>sport
  tmp=${tmp//dst/src}
  tmp=${tmp//dport/sport}
  # fulltext replace src=>dst, sport=>dport
  tmp=${tmp//sr_c/dst}
  tmp=${tmp//spor_t/dport}
  str_r=${tmp}

  # setup rule
    tc filter add dev imq0 protocol ip parent 1:0 prio $str_p u32 \
        $str_m flowid 1:${CLSID_INET}${str_q}
    tc filter add dev imq1 protocol ip parent 2:0 prio $str_p u32 \
        $str_r flowid 2:${CLSID_INET}${str_q}

  # if no "lan" or "wan" keywords were used
  # => doublicate the rule for another direction
  if test $flag_both -eq 1 -a "$str_m" != "$str_r"; then
    tc filter add dev imq0 protocol ip parent 1:0 prio $str_p u32 \
        $str_r flowid 1:${CLSID_INET}${str_q}
    tc filter add dev imq1 protocol ip parent 2:0 prio $str_p u32 \
        $str_m flowid 2:${CLSID_INET}${str_q}
  fi
}

#############################################################################
#                          Zone filters setup                               #
#############################################################################

echo -n "Setup zones"
i=0
for fp in $WAN_ZONES; do
if test "$fp" != "inet"; then
  fname="/opt/etc/nshaper/ip_$fp.lst"
  if test -e "$fname"; then
    echo -n " $fp";
    # read ip range from file
    while read ADDR B; do
      if test -n "$ADDR"; then
        if test `expr "$ADDR" : '\(.\)'` != "#"; then
          #echo "Adding $ADDR"
          tc filter add dev imq0 protocol ip parent 1:0 prio 1 u32 \
            match ip src `toprefixmask $ADDR` \
            flowid 1:2${i}
          tc filter add dev imq1 protocol ip parent 2:0 prio 1 u32 \
            match ip dst `toprefixmask $ADDR` \
            flowid 2:2${i}
        fi
      fi
    done < $fname
  fi
fi
i=$(($i+1))
done
echo ""

#############################################################################
#                    Application/user filters setup                         #
#############################################################################

echo "Applying rules "

# set rules for time-critical queue (both users' and router's traffic)
setrule proto icmp prio 10 queue 0		# ICMP
setrule port 53 queue 0				# DNS
setrule short flag ack queue 0			# TCP ACK
setrule short flag syn queue 0			# TCP SYN
setrule short flag rst queue 0			# TCP RST


# extract users' traffic using destiation ip
LAN="$LAN_IP/$LAN_MASK"

# set rules for high-priority queue
setrule lan ip $LAN afc 4 queue 1		# assured forward GOLD class
						# (af4x obsoletes TOS=10)
setrule lan ip $LAN wan port 554 queue 1	# MMS/RealMedia
setrule lan ip $LAN wan port 1755 queue 1	# MMS
setrule lan ip $LAN wan port 1935 queue 1	# RTMP/Flash

# set rules for middle-priority queue
setrule lan ip $LAN afc 3 queue 2		# assured forward SILVER class
setrule lan ip $LAN wan port 80 queue 2		# HTTP
setrule lan ip $LAN wan port 443 queue 2	# HTTPS
setrule lan ip $LAN wan port 5190 queue 2	# AIM

# set rules for low-prority queue
setrule lan ip $LAN queue 3			# all remaining users' traffic 


# only router's traffic remains unfiltered here
setrule lan port 50022 queue 1  		# router's SSH
setrule lan port 8081  queue 1  		# router's WEB server

# move all other router's traffic (such as P2P) to lazy queue
setrule lan all queue 4				# all other traffic to router



# finally route incoming/outgoing traffic to imq0/imq1 interfaces
iptables -t mangle -I PREROUTING -i $WAN_IF -j IMQ --todev 0 > /dev/null 2>&1
iptables -t mangle -A POSTROUTING -o $WAN_IF -j IMQ --todev 1 > /dev/null 2>&1

}

# main
case "$1" in
stop)	stop
  	echo "nShaper stopped.";;
start)	
	clear_cfg
	start
	echo "nShaper started.";;
restart)
	clear_cfg
	start
	echo "nShaper restarted.";;
*)	echo ""
  	echo "nShaper - upgrade your Internet!"
  	echo ""
  	echo "Command line: $0 { start | stop | restart }";;
esac

return 0
