limitator/limitator

#!/bin/sh /etc/rc.common
START=95
STOP=01

EXTRA_COMMANDS="save reset show json"
EXTRA_HELP=<<EOF
        reset Resets quota
        save Saves counters to persistent storage
        show Shows remaining quota
        json Print data in JSON format
EOF

config_load "limitator"
config_get MONTHLY_QUOTA limitator quota 15 # Monthly quota in GB
config_get WAN_DEV limitator wan_dev 'wlan0' # Interface name for WAN
config_get FIRST_DAY limitator first_day 13 # Day of month when internet plan starts accounting..
config_get START_TIME limitator start_time '08:00' # Starting time of internet plan limited hours.
config_get END_TIME limitator end_time '22:00' # Ending time of internet plan limited hours.
config_get INTRANET limitator intranet '10.0.0.0/8' # Traffic in Intranet does not count for quota.
config_get EXCLUDE_NETWORK limitator exclude_network # Traffic originated in specified network does not count for quota.
config_get EXCLUDE_NETWORK2 limitator exclude_network2 # Traffic originated in specified network does not count for quota.
config_get EXCLUDE_NETWORK3 limitator exclude_network3 # Traffic originated in specified network does not count for quota.
config_get USED_BYTES_FILE limitator used_bytes_file '/root/limitator/used_bytes' # A persistent file to save used bytes value.
config_get LOG_FILE limitator log_file '/root/limitator/log' # A persistent file to log.
config_get IPTABLES_CHAIN limitator iptables_chain 'LIMITATOR' # Name of iptables chain.
config_get IPTABLES_FORWARD limitator iptables_forward 'forwarding_rule' # Name of iptables rule for forward.
config_get IPTABLES_INPUT limitator iptables_input 'input_rule' # Name of iptables rule for input.
config_get IPTABLES_OUTPUT limitator iptables_output 'output_rule' # Name of iptables rule for output.

get_saved_used_bytes(){
    local used_bytes
    if [ -s "$USED_BYTES_FILE" ];then
        used_bytes=$(cat $USED_BYTES_FILE 2>/dev/null)
    else
        used_bytes=0
    fi
    echo -n $used_bytes
}

get_iptables_quota_rule_num(){
    if [ -n "$EXCLUDE_NETWORK" ];then
        if [ -n "$EXCLUDE_NETWORK2" ];then
            if [ -n "$EXCLUDE_NETWORK3" ];then
                echo -n 7
            else
                echo -n 5
            fi
        else
            echo -n 3
        fi
    else
        echo -n 1
    fi
}

get_current_used_bytes(){
    local used_bytes=$(iptables --wait -vnxL $IPTABLES_CHAIN $(get_iptables_quota_rule_num) | awk '{print $2}')
    if [ -n "$used_bytes" ];then
        echo -n $used_bytes
    else
        get_saved_used_bytes
    fi
}

set_saved_used_bytes(){
    local used_bytes=$(get_current_used_bytes)
    if [ $used_bytes -gt 0 ];then
        echo -n $used_bytes > $USED_BYTES_FILE
    fi
}

get_total_quota_until_today(){
    local current_time=$(date +'%s')
    local first_day=$(date -d $(date +"%Y-%m-$FIRST_DAY") +'%s')
    local next_first_day
   
    if [ $current_time -ge $first_day ]; then
        next_first_day=$(date -d $(date -d @$first_day +%Y-%m-%d)"+1 month" +'%s')
    else
        next_first_day=$first_day
        first_day=$(date -d $(date -d "1 month ago" +"%Y-%m-$FIRST_DAY") +'%s')
    fi
  
    local seconds_in_month=$(($next_first_day - $first_day))
    local montly_quota_bytes=$(($MONTHLY_QUOTA * 1024 * 1024 * 1024))
    local elapsed_seconds_tomorrow=$(($(date -d $(date -d "1 day" +'%Y-%m-%d') +'%s') - $first_day))
    
    echo -n $(($montly_quota_bytes * $elapsed_seconds_tomorrow / $seconds_in_month))
}

get_remaining_quota(){
    echo -n $(($(get_total_quota_until_today) - $(get_current_used_bytes)))
}

show_remaining_quota(){
    echo Remaining: $(($(get_remaining_quota) / 1024 / 1024)) MB
    echo Total quota until today: $(($(get_total_quota_until_today) / 1024 / 1024)) MB
}

get_json(){
    echo -n "{ \"remaining\": $(get_remaining_quota), \"used_bytes\": $(get_current_used_bytes), \"current_quota\": $(get_total_quota_until_today), \"total_quota\": $(($MONTHLY_QUOTA * 1024 * 1024 * 1024)) }"
}

reset_quota(){
    iptables --wait -Z $IPTABLES_CHAIN $(get_iptables_quota_rule_num)
    echo -n 0 > ${USED_BYTES_FILE}
}

del_iptables_rule(){
    iptables --wait -F $IPTABLES_CHAIN
    for chain_name in $IPTABLES_FORWARD $IPTABLES_INPUT $IPTABLES_OUTPUT
    do
        while true
        do
            line=$(iptables --wait -L $chain_name|sed -e '1,2d'|grep -n $IPTABLES_CHAIN|head -1|cut -d : -f 1 2>/dev/null)
            [ -n "$(echo $line|grep -E ^[[:digit:]]+$)" ] && iptables --wait -D $chain_name $line || break
        done
    done
    iptables --wait -D $IPTABLES_FORWARD 1
    iptables --wait -X $IPTABLES_CHAIN
}

set_iptables_rule(){
    iptables --wait -N $IPTABLES_CHAIN
    if [ -n "$EXCLUDE_NETWORK" ];then
        iptables --wait -I $IPTABLES_CHAIN 1 -s $EXCLUDE_NETWORK -j ACCEPT
        iptables --wait -I $IPTABLES_CHAIN 2 -d $EXCLUDE_NETWORK -j ACCEPT
    fi
    if [ -n "$EXCLUDE_NETWORK2" ];then
        iptables --wait -I $IPTABLES_CHAIN 3 -s $EXCLUDE_NETWORK2 -j ACCEPT
        iptables --wait -I $IPTABLES_CHAIN 4 -d $EXCLUDE_NETWORK2 -j ACCEPT
    fi
    if [ -n "$EXCLUDE_NETWORK3" ];then
        iptables --wait -I $IPTABLES_CHAIN 5 -s $EXCLUDE_NETWORK3 -j ACCEPT
        iptables --wait -I $IPTABLES_CHAIN 6 -d $EXCLUDE_NETWORK3 -j ACCEPT
    fi
    iptables --wait -I $IPTABLES_CHAIN $(get_iptables_quota_rule_num) -c 0 $(get_saved_used_bytes) -m quota --quota $(get_remaining_quota) -j ACCEPT 2>> $LOG_FILE
    iptables --wait -I $IPTABLES_CHAIN $(($(get_iptables_quota_rule_num) + 1)) -j REJECT 2>> $LOG_FILE

    # Iptables requires UTC time
    local start_time_utc=$(date -u +%H:%M --date=@$(date -d "today $START_TIME" +%s))
    local end_time_utc=$(date -u +%H:%M --date=@$(date -d "today $END_TIME" +%s))

    
    iptables --wait -I $IPTABLES_FORWARD 1 -d $INTRANET -s $INTRANET -j ACCEPT
    iptables -I $IPTABLES_FORWARD 2 -m time --timestart $start_time_utc --timestop $end_time_utc -j $IPTABLES_CHAIN
    iptables -I $IPTABLES_INPUT -i $WAN_DEV ! -s $INTRANET -m time --timestart $start_time_utc --timestop $end_time_utc -j $IPTABLES_CHAIN
    iptables -I $IPTABLES_OUTPUT -o $WAN_DEV ! -d $INTRANET -m time --timestart $start_time_utc --timestop $end_time_utc -j $IPTABLES_CHAIN
}

start() {
    set_iptables_rule
}

stop() {
    set_saved_used_bytes
    echo $(date +'%F %H:%M') used_bytes=$(get_current_used_bytes) >> $LOG_FILE
    del_iptables_rule
}

restart() {
    stop
    start
}

save() {
    set_saved_used_bytes
}

reset() {
    reset_quota
    restart
}

show() {
    show_remaining_quota
}

json() {
    get_json
}