#!/bin/sh

. /lib/functions.sh
. /usr/share/libubox/jshn.sh
. /lib/delos-functions.sh

PACKAGE="station_quota"
VERSION="v2.0.0"

# the lock is meant for ebtables, and other scripts are also
# using ebtables. Needs clean up, maybe move to functions.sh, TBD?
LOCKFILE=/tmp/ebtables.lock

EVENTFILE=/tmp/station_quota_events.txt
CONFIG_SYNC_RETRIES=3
CONFIG_SYNC_DUPLICATE_TIME=2
CONFIG_SYNC_KEEP_ENTRIES=100

# ubus timer namespace prefix
TIMER_PREFIX=quota-

##### logging

function logger() {
    if [ $CFG_LOGGING -eq 1 ]; then
        command logger -t "$PACKAGE" "$@"
    fi
}

##### ebtables handling

ebt_create_chains() {
    ebtables -N "i-${PACKAGE}" -P RETURN 2>/dev/null || return 0
    ebtables -N "o-${PACKAGE}" -P RETURN 2>/dev/null || return 0
    ebtables -A INPUT   -i ath+ -j "i-${PACKAGE}"
    ebtables -A FORWARD -i ath+ -j "i-${PACKAGE}"
    ebtables -A FORWARD -o ath+ -j "o-${PACKAGE}"
    ebtables -A OUTPUT  -o ath+ -j "o-${PACKAGE}"
}

ebt_delete_chains() {
    ebtables -D INPUT   -i ath+ -j "i-${PACKAGE}" 2>/dev/null
    ebtables -D FORWARD -i ath+ -j "i-${PACKAGE}" 2>/dev/null
    ebtables -D FORWARD -o ath+ -j "o-${PACKAGE}" 2>/dev/null
    ebtables -D OUTPUT  -o ath+ -j "o-${PACKAGE}" 2>/dev/null
    ebtables -X "i-${PACKAGE}" 2>/dev/null
    ebtables -X "o-${PACKAGE}" 2>/dev/null
}

ebt_add_macs_to_blacklist() {
    local macs="$1"
    
    local mac
    for mac in $macs; do
        logger "Adding $mac to blacklist"
        ebtables -D "i-${PACKAGE}" -s "$mac" -j DROP 2>/dev/null
        ebtables -A "i-${PACKAGE}" -s "$mac" -j DROP
        ebtables -D "o-${PACKAGE}" -d "$mac" -j DROP 2>/dev/null
        ebtables -A "o-${PACKAGE}" -d "$mac" -j DROP
    done
}

ebt_remove_macs_from_blacklist() {
    local macs="$1"
    
    local mac
    for mac in $macs; do
        logger "Removing $mac from blacklist"
        ebtables -D "i-${PACKAGE}" -s "$mac" -j DROP 2>/dev/null
        ebtables -D "o-${PACKAGE}" -d "$mac" -j DROP 2>/dev/null
    done
}

##### timer handling

timer_to_mac() {
    local timer_id="$1"
    
    echo "${timer_id#${TIMER_PREFIX}}"
}

mac_to_timer() {
    local mac="$1"
    
    echo "${TIMER_PREFIX}${mac}"
}

# High-level code should not call this directly. Call quota_operation instead in order to
# ensure proper blacklist handling.
timer_operation() {
    local operation="$1"
    local mac="$2"
    local arg1="$3"
    local arg2="$4"
    local arg3="$5"

    local timer_id=$(mac_to_timer "$mac")
    
    # logger "Timer operation '$operation' for '$mac' (arg1='$arg1', arg2='$arg2', arg3='$arg3')"
    
    case "$operation" in
        list)
            echo $(ubus call timer "$operation" "{\"id\":\"${timer_id}\"}" 2>/dev/null | jsonfilter -e "@.ids[@~\"^${TIMER_PREFIX}\"]")
            ;;
        query)
            echo $(ubus call timer "$operation" "{\"id\":\"${timer_id}\"}" 2>/dev/null)
            ;;
        getvar)
            local var="$arg1"
            echo $(ubus call timer query "{\"id\":\"${timer_id}\"}" 2>/dev/null | jsonfilter -e '$.'"$var" 2>/dev/null)
            ;;
        pause|resume|remove)
            ubus call timer "$operation" "{\"id\":\"${timer_id}\"}" 2>/dev/null
            ;;
        new)
            ubus call timer remove "{\"id\":\"${timer_id}\"}" 2>/dev/null

            local timeout="${arg1:-0}"
            local remaining="${arg2:-${timeout}}"
            local start=$(( $timeout - $remaining ))
            [ $start -lt 0 ] && start=0
            
            ubus call timer add "{\"id\":\"${timer_id}\",\"command\":\"${MYSELF}\",\"timeout\":${timeout},\"start\":${start}}"
            ;;
        *)
            logger "Unknown timer operation '$operation'"
            ;;
    esac
}

##### get configured quota

config_get_quota_cb() {
    local name="$1"
    local mac="$2"
    local day="$3"
    local option="$4"
    local result_var="$5"

    local entry_mac=$(normalize_mac $(config_get "$name" "station"))
    local entry_option_val=$(config_get "$name" "$option")
    local entry_days=$(config_get "$name" "daysofweek")

    case "$entry_days" in
        *${day}*)
            [ "$entry_mac" = "$mac" ] && append "$result_var" "$entry_option_val"
            ;;
    esac
}

config_get_quota() {
    local mac="$1"

    local result=""
    local day="$(date +%A)"

    config_foreach config_get_quota_cb "entry" "$mac" "$day" "quota" "result"
    
    echo "$result"
}

config_list_stations_cb() {
    local name="$1"
    local result_var="$2"

    local entry_mac=$(normalize_mac $(config_get "$name" "station"))
    local previous_macs=$(eval echo -n \$$result_var)

    if [ -z $(echo "$previous_macs" | grep "$entry_mac") ]; then
        append "$result_var" "$entry_mac"
    fi
}

config_list_stations() {
    local result=""

    config_foreach config_list_stations_cb "entry" "result"
    
    echo "$result"
}

##### WiFi helpers

wifi_get_connected_stations() {
    local result=""

    local devices
    json_load "$(ubus call iwinfo devices)"
    json_get_values devices devices
    
    local device
    for device in $devices; do
        local clients=""
        eval $(ubus call "hostapd.$device" get_clients 2>/dev/null | jsonfilter -e 'clients=$.clients' 2>/dev/null)
        logger "Connected stations @ $device: $clients"
        
        local client
        for client in $clients; do
            append result $(normalize_mac "$client")
        done
    done
    
    echo "$result"
}

wifi_is_station_connected() {
    local mac="$1"
    local stations=$(wifi_get_connected_stations)
    echo "$stations" | grep "$mac"
}

wifi_kick_station() {
    local mac="$1"
    
    local vap
    for vap in $(ubus call iwinfo devices | jsonfilter -e '$.devices[@~"^ath"]'); do
        iwpriv "$vap" kickmac "$mac"
    done
}

##### misc. helpers

normalize_mac() {
    local mac="$1"
    
    echo "$mac" | awk '{print toupper($0)}'
}

##### config sync

config_sync_send_event() {
    local event="$1"
    
    logger "Sending sync event '$event'"
    for I in $(seq 1 $CONFIG_SYNC_RETRIES); do
        ubus call configsync send_message "{\"service\":\"quota\",\"event\":\"${event}\"}"
    done
}

config_sync_send_station_event() {
    local event="$1"
    local mac="$2"
    
    local remaining=$(quota_operation getvar "$mac" remaining)
    if [ -n "$remaining" ]; then
        logger "Sending sync event '$event' for '$mac': remaining $remaining"
        for I in $(seq 1 $CONFIG_SYNC_RETRIES); do
            ubus call configsync send_message "{\"service\":\"quota\",\"event\":\"${event}\",\"mac\":\"${mac}\",\"remaining\":${remaining}}"
        done
    fi
}

config_sync_write_event() {
    echo "$(date +%s) '$1' '$2' '$3'" >> "$EVENTFILE"
    
    tail -n $CONFIG_SYNC_KEEP_ENTRIES "$EVENTFILE" > "$EVENTFILE".tmp
    mv -f "$EVENTFILE".tmp "$EVENTFILE"
}

config_sync_is_duplicate_event() {
    local difftime=""
    
    [ -f "$EVENTFILE" ] || return 1
    local line=$(cat "$EVENTFILE" | grep " '$1' '$2' '$3'\$" | tail -n 1)
    if [ -n "$line" ]; then
        local evtime=$(echo "$line" | sed -e 's/ .*//')
        local nowtime=$(date +%s)
        difftime=$(( $nowtime - $evtime ))
    fi
    
    if [ -n "$difftime" ] && [ $difftime -le $CONFIG_SYNC_DUPLICATE_TIME ]; then
        logger "Ignoring duplicate event"
        return 0
    else
        return 1
    fi
}

##### station quota program logic

update_blacklist_entry() {
    local mac="$1"
    
    local cfg_quota=$(config_get_quota "$mac")
    local state=$(quota_operation getvar "$mac" state)

    if [ -z "$cfg_quota" ]; then
        ebt_remove_macs_from_blacklist "$mac"
    elif [ -z "$state" ] || [ "$state" = "RUNNING" ]; then
        ebt_remove_macs_from_blacklist "$mac"
    else
        ebt_add_macs_to_blacklist "$mac"  # paused or expired timer
    fi
}

quota_operation() {
    local operation="$1"
    local mac="$2"
    local arg1="$3"
    local arg2="$4"
    local arg3="$5"

    # logger "Quota operation '$operation' for '$mac' (arg1='$arg1', arg2='$arg2', arg3='$arg3')"
    
    # operate timer
    case "$operation" in
        list|query|getvar)
            echo $(timer_operation "$operation" "$mac" "$arg1" "$arg2" "$arg3")
            ;;
        pause|resume|new|remove)
            timer_operation "$operation" "$mac" "$arg1" "$arg2" "$arg3"
            ;;
        updatebl|*)
            ;;
    esac

    # operate blacklist
    case "$operation" in
        list|query|getvar)
            ;;
        pause|resume|new|updatebl)
            update_blacklist_entry "$mac"
            ;;
        remove)
            ebt_remove_macs_from_blacklist "$mac"
            ;;
        *)
            logger "Unknown quota operation '$operation'"
            ;;
    esac
}

station_associated() {
    local mac="$1"
    local sync="${2:-1}"
    
    logger "Station $mac associated (sync=$sync)"
    local state=$(quota_operation getvar "$mac" state)
    if [ "$state" != "EXPIRED" ]; then
        if [ -z "$state" ]; then
            local cfg_quota=$(config_get_quota "$mac")
            if [ -n "$cfg_quota" ]; then
                logger "Starting new quota for $mac ($cfg_quota)"
                quota_operation new "$mac" $(dvl_time_to_mseconds "$cfg_quota")
            fi
        else
            logger "Resuming quota for $mac"
            quota_operation resume "$mac"
        fi
        [ $sync -eq 1 ] && config_sync_send_station_event assoc "$mac"
    fi
}

station_disassociated() {
    local mac="$1"
    local sync="${2:-1}"

    logger "Station $mac disassociated (sync=$sync)"
    local state=$(quota_operation getvar "$mac" state)
    if [ "$state" != "EXPIRED" ]; then
        if [ -n "$state" ]; then
            quota_operation pause "$mac"
            [ $sync -eq 1 ] && config_sync_send_station_event disassoc "$mac"
        fi
    fi
}

station_timer_expired() {
    local mac="$1"
    
    logger "Timer for station $mac expired"
    local state=$(quota_operation getvar "$mac" state)
    if [ -n "$state" ]; then
        quota_operation updatebl "$mac"
        wifi_kick_station "$mac"
    fi
}

station_updated_by_remote() {
    local mac="$1"
    local remaining="$2"

    logger "Quota for station $mac updated by remote (remaining $remaining)"

    local mseconds="$remaining"

    local cfg_quota=$(config_get_quota "$mac")
    if [ -n "$cfg_quota" ]; then
        mseconds=$(dvl_time_to_mseconds "$cfg_quota")
    fi

    [ $remaining -le 0 ] && remaining=1  # timer wouldn't get created with zero time remaining
    
    quota_operation new "$mac" "$mseconds" "$remaining"
}

midnight() {
    for timer in $(quota_operation list); do
        local mac=$(timer_to_mac "$timer")
        local cfg_quota=$(config_get_quota "$mac")
        local state=$(quota_operation getvar "$mac" state)

        if [ -n "$cfg_quota" ] && [ "$state" = "RUNNING" ]; then
            logger "Refreshing running quota for station $mac ($cfg_quota)"
            quota_operation new "$mac" $(dvl_time_to_mseconds "$cfg_quota")
        else
            logger "Removing stale quota for station $mac"
            quota_operation remove "$mac"
        fi
    done
}

request_all_quotas_from_remote() {
    logger "Requesting quota from remote devices"
    config_sync_send_event list
}

announce_all_quotas_to_remote() {
    for timer in $(quota_operation list); do
        local mac=$(timer_to_mac "$timer")
        local state=$(quota_operation getvar "$mac" state)
        local remaining=$(quota_operation getvar "$mac" remaining)
        
        if [ -n "$state" ]; then
            local event=""
            case "$state" in
                RUNNING)
                    event="assoc"
                    ;;
                PAUSED|EXPIRED)
                    event="disassoc"
                    ;;
            esac

            if [ -n "$event" ]; then            
                logger "Announcing station $mac ($event, remaining $remaining)"
                config_sync_send_station_event "$event" "$mac" "$remaining"
            fi
        fi
    done
}

update_quotas_for_connected_stations() {
    for station in $(wifi_get_connected_stations); do
        logger "Adding quota for connected station $mac"
        station_associated "$station" 1
    done
}

remove_all_station_quotas() {
    for timer in $(quota_operation list); do
        quota_operation remove $(timer_to_mac "$timer")
    done
}

reload_quota_config() {
    # remove quotas for stations which have no configured quota anymore
    for timer in $(quota_operation list); do
        local mac=$(timer_to_mac "$timer")
        local cfg_quota=$(config_get_quota "$mac")

        if [ -z "$cfg_quota" ]; then
            logger "Removing stale quota for station $mac"
            quota_operation remove "$mac"
        fi
    done

    # update quotas for configured stations
    for mac in $(config_list_stations); do
        local cfg_quota=$(config_get_quota "$mac")
        if [ -n "$cfg_quota" ]; then
            local cfg_mseconds=$(dvl_time_to_mseconds "$cfg_quota")
            local state=$(quota_operation getvar "$mac" state)
            local timeout=$(quota_operation getvar "$mac" timeout)
            
            if [ -z "$timeout" ] || [ "$timeout" -ne "$cfg_mseconds" ]; then
                # configured quota is new or different from timer timeout
                if [ "$state" = "RUNNING" ] || wifi_is_station_connected "$mac"; then
                    logger "Changing quota for connected station $mac ($cfg_quota)"
                    quota_operation new "$mac" "$cfg_mseconds"
                    config_sync_send_station_event assoc "$mac" "$cfg_mseconds"
                else
                    logger "Removing quota for disconnected station $mac ($cfg_quota)"
                    quota_operation remove "$mac"
                fi
            fi
        fi
    done
}

##### service handling

start_service() {
    ebt_create_chains
    crontab -l | grep -v "$MYSELF" | crontab -
    (crontab -l; echo "00 00 * * * $MYSELF --event midnight") | sort -u | crontab -
    
    > "$EVENTFILE"
    remove_all_station_quotas
    (sleep 120; "$MYSELF" -e request_list; sleep 10; "$MYSELF" -e update_connected) &
}

stop_service() {
    rm -f "$EVENTFILE"
    remove_all_station_quotas

    crontab -l | grep -v "$MYSELF" | crontab -
    ebt_delete_chains
}

##### option handling

get_program_options() {
    OPT_EVENT=""
    OPT_MAC=""
    OPT_REMAINING=""
    OPT_TIMER_ID=""
    
    if [ -n "$TIMER_ID" ]; then
        # called by timer
        OPT_EVENT="timer_expired"
        OPT_TIMER_ID="$TIMER_ID"
        OPT_MAC=$(normalize_mac $(timer_to_mac "$TIMER_ID"))
    else
        # called by script or cron
        while [ $# -gt 0 ]; do
            local key="$1"
            shift
            case "$key" in
                start|stop)
                    OPT_EVENT="$key"
                    ;;
                -e|--event)
                    OPT_EVENT="$1"
                    shift
                    ;;
                -m|--mac)
                    OPT_MAC=$(normalize_mac "$1")
                    shift
                    ;;
                -r|--remaining)
                    OPT_REMAINING="$1"
                    shift
                    ;;
                *)
                    logger "Unknown command line option '$key'"
                    ;;
            esac
        done
    fi
}

##### main

config_load "$PACKAGE"
config_get_bool CFG_PACKAGE_ENABLED global enabled 0
config_get_bool CFG_LOGGING global logging 0

MYSELF="$0"
RETURN_VALUE=0

lock "$LOCKFILE"

if [ $CFG_PACKAGE_ENABLED -eq 1 ]; then
    get_program_options "$@"

    logger "Event '$OPT_EVENT' for station '$OPT_MAC', remaining '$OPT_REMAINING'"
    
    case "$OPT_EVENT" in
        start)
            start_service
            ;;
        stop)
            stop_service
            ;;
        reload)
            initialized=$(ebtables -L | grep "$PACKAGE")
            if [ -z "$initialized" ]; then
                logger "Service $PACKAGE was not started, initializing"
                start_service
            fi
            reload_quota_config
            ;;
        assoc)
            station_associated "$OPT_MAC" 1
            ;;
        disassoc)
            station_disassociated "$OPT_MAC" 1
            ;;
        request_list)
            request_all_quotas_from_remote
            ;;
        midnight)
            midnight
            ;;
        timer_expired)
            station_timer_expired "$OPT_MAC"
            ;;
        update_connected)
            update_quotas_for_connected_stations
            ;;
        assoc_remote)
            if [ -n "$OPT_REMAINING" ] && ! config_sync_is_duplicate_event "$OPT_EVENT" "$OPT_MAC" "$OPT_REMAINING"; then
                config_sync_write_event "$OPT_EVENT" "$OPT_MAC" "$OPT_REMAINING"
                station_updated_by_remote "$OPT_MAC" "$OPT_REMAINING"
                station_associated "$OPT_MAC" 0
            fi
            ;;
        disassoc_remote)
            if [ -n "$OPT_REMAINING" ] && ! config_sync_is_duplicate_event "$OPT_EVENT" "$OPT_MAC" "$OPT_REMAINING"; then
                config_sync_write_event "$OPT_EVENT" "$OPT_MAC" "$OPT_REMAINING"
                station_updated_by_remote "$OPT_MAC" "$OPT_REMAINING"
                station_disassociated "$OPT_MAC" 0
            fi
            ;;
        list_remote)
            if ! config_sync_is_duplicate_event "$OPT_EVENT" "$OPT_MAC" "$OPT_REMAINING"; then
                config_sync_write_event "$OPT_EVENT"
                announce_all_quotas_to_remote
            fi
            ;;
        test)
            ;;
        *)
            logger "Unknown event '$OPT_EVENT'"
            RETURN_VALUE=1
            ;;
    esac
else
    initialized=$(ebtables -L | grep "$PACKAGE")
    if [ -n "$initialized" ]; then
        logger "Service $PACKAGE was running, stopping"
        stop_service
    fi
    logger "Feature $PACKAGE disabled"
fi

lock -u "$LOCKFILE"

exit $RETURN_VALUE