#!/bin/sh
#
# Copyright (C) 2017 openwrt-ssr
# Copyright (C) 2017 yushi studio <ywb94@qq.com>
#
# This is free software, licensed under the GNU General Public License v3.
# See /LICENSE for more information.
#

TAG="_SS_SPEC_RULE_"                          # comment tag
IPT="iptables -t nat"                         # alias of iptables
FWI=$(uci get firewall.vssr.path 2>/dev/null) # firewall include file

usage() {
    cat <<-EOF
		Usage: vssr-rules [options]

		Valid options are:

		    -s <server_ip>          ip address of vssr remote server
		    -l <local_port>         port number of vssr local server
		    -S <server_ip>          ip address of vssr remote UDP server
		    -L <local_port>         port number of vssr local UDP server
		    -i <ip_list_file>       a file content is bypassed ip list
		    -a <lan_ips>            lan ip of access control, need a prefix to
		                            define access control mode
		    -b <wan_ips>            wan ip of will be bypassed
		    -w <wan_ips>            wan ip of will be forwarded
		    -p <fp_lan_ips>         lan ip of will be global proxy
		    -G <gm_lan_ips>         lan ip of will be game mode proxy
		    -D <proxy_ports>        proxy ports
		    -e <extra_options>      extra options for iptables
		    -o                      apply the rules to the OUTPUT chain
		    -O                      apply the global rules to the OUTPUT chain
		    -u                      enable udprelay mode, TPROXY is required
		    -U                      enable udprelay mode, using different IP
		                            and ports for TCP and UDP
		    -f                      flush the rules
		    -g                      gfw list mode
		    -r                      return china mode
		    -h                      show this help message and exit
	EOF
    exit $1
}

loger() {
    # 1.alert 2.crit 3.err 4.warn 5.notice 6.info 7.debug
    logger -st vssr-rules[$$] -p$1 $2
}

flush_r() {
    flush_iptables() {
        local ipt="iptables -t $1"
        local DAT=$(iptables-save -t $1)
        eval $(echo "$DAT" | grep "$TAG" | sed -e 's/^-A/$ipt -D/' -e 's/$/;/')
        for chain in $(echo "$DAT" | awk '/^:SS_SPEC/{print $1}'); do
            $ipt -F ${chain:1} 2>/dev/null && $ipt -X ${chain:1}
        done
    }
    flush_iptables nat
    flush_iptables mangle
    ip rule del fwmark 0x01/0x01 table 100 2>/dev/null
    ip route del local 0.0.0.0/0 dev lo table 100 2>/dev/null
    ipset -X ss_spec_lan_ac 2>/dev/null
    ipset -X ss_spec_wan_ac 2>/dev/null
    ipset -X ssr_gen_router 2>/dev/null
    ipset -X fplan 2>/dev/null
    ipset -X gmlan 2>/dev/null
    ipset -X oversea 2>/dev/null
    ipset -X whitelist 2>/dev/null
    ipset -X blacklist 2>/dev/null
    ipset -X gfwlist 2>/dev/null
    [ -n "$FWI" ] && echo '#!/bin/sh' >$FWI
    return 0
}

ipset_r() {
    ipset -N gmlan hash:net hashsize 4096 maxelem 1000000 2>/dev/null
    $IPT -N SS_SPEC_WAN_AC
    $IPT -I SS_SPEC_WAN_AC -p tcp ! --dport 53 -d $server -j RETURN
    for ip in $LAN_GM_IP; do ipset -! add gmlan $ip; done
    case "$RUNMODE" in
    router)
        ipset -! -R <<-EOF || return 1
			create ss_spec_wan_ac hash:net hashsize 4096 maxelem 1000000
			$(gen_iplist | sed -e "s/^/add ss_spec_wan_ac /")
		EOF
        ipset -N gfwlist hash:net hashsize 4096 maxelem 1000000 2>/dev/null
        $IPT -A SS_SPEC_WAN_AC -m set --match-set ss_spec_wan_ac dst -j RETURN
		$IPT -A SS_SPEC_WAN_AC -m set --match-set gmlan src -m set ! --match-set china dst -j SS_SPEC_WAN_FW
		$IPT -A SS_SPEC_WAN_AC -m set --match-set china dst -j RETURN
		$IPT -A SS_SPEC_WAN_AC -j SS_SPEC_WAN_FW
        ;;
    gfw|direct)
        ipset -N gfwlist hash:net hashsize 4096 maxelem 1000000 2>/dev/null
        $IPT -A SS_SPEC_WAN_AC -m set --match-set gfwlist dst -j SS_SPEC_WAN_FW
        $IPT -A SS_SPEC_WAN_AC -m set --match-set gmlan src -m set ! --match-set china dst -j SS_SPEC_WAN_FW
        $IPT -A SS_SPEC_WAN_AC -m set --match-set china dst -j RETURN
        ;;
    oversea)
        ipset -N oversea hash:net hashsize 4096 maxelem 1000000 2>/dev/null
        $IPT -I SS_SPEC_WAN_AC -m set --match-set oversea dst -j SS_SPEC_WAN_FW
        $IPT -A SS_SPEC_WAN_AC -m set --match-set gmlan src -j SS_SPEC_WAN_FW
        $IPT -A SS_SPEC_WAN_AC -m set --match-set china dst -j SS_SPEC_WAN_FW
        ;;
    all)
        $IPT -A SS_SPEC_WAN_AC -j SS_SPEC_WAN_FW
        ;;
    esac

    ipset -N fplan hash:net hashsize 4096 maxelem 1000000 2>/dev/null
    for ip in $LAN_FP_IP; do ipset -! add fplan $ip; done
    $IPT -I SS_SPEC_WAN_AC -m set --match-set fplan src -j SS_SPEC_WAN_FW

    ipset -N whitelist hash:net hashsize 4096 maxelem 1000000 2>/dev/null
    ipset -N blacklist hash:net hashsize 4096 maxelem 1000000 2>/dev/null
    $IPT -I SS_SPEC_WAN_AC -m set --match-set blacklist dst -j SS_SPEC_WAN_FW
    $IPT -I SS_SPEC_WAN_AC -m set --match-set whitelist dst -j RETURN

    for ip in $WAN_BP_IP; do ipset -! add whitelist $ip; done
    for ip in $WAN_FW_IP; do ipset -! add blacklist $ip; done

    return $?
}

fw_rule() {
    $IPT -N SS_SPEC_WAN_FW
    $IPT -A SS_SPEC_WAN_FW -d 0.0.0.0/8 -j RETURN
    $IPT -A SS_SPEC_WAN_FW -d 10.0.0.0/8 -j RETURN
    $IPT -A SS_SPEC_WAN_FW -d 127.0.0.0/8 -j RETURN
    $IPT -A SS_SPEC_WAN_FW -d 169.254.0.0/16 -j RETURN
    $IPT -A SS_SPEC_WAN_FW -d 172.16.0.0/12 -j RETURN
    $IPT -A SS_SPEC_WAN_FW -d 192.168.0.0/16 -j RETURN
    $IPT -A SS_SPEC_WAN_FW -d 224.0.0.0/4 -j RETURN
    $IPT -A SS_SPEC_WAN_FW -d 240.0.0.0/4 -j RETURN
    $IPT -A SS_SPEC_WAN_FW -p tcp $PROXY_PORTS \
        -j REDIRECT --to-ports $local_port 2>/dev/null || {
        loger 3 "Can't redirect, please check the iptables."
        exit 1
    }
    return $?
}

ac_rule() {
    if [ -n "$LAN_AC_IP" ]; then
        case "${LAN_AC_IP:0:1}" in
        w | W)
            MATCH_SET="-m set --match-set ss_spec_lan_ac src"
            ;;
        b | B)
            MATCH_SET="-m set ! --match-set ss_spec_lan_ac src"
            ;;
        *)
            loger 3 "Bad argument \`-a $LAN_AC_IP\`."
            return 2
            ;;
        esac
    fi
    IFNAME=$(uci get -P/var/state network.lan.device 2>/dev/null)
    ipset -! -R <<-EOF || return 1
		create ss_spec_lan_ac hash:net hashsize 4096 maxelem 1000000
		$(for ip in ${LAN_AC_IP:1}; do echo "add ss_spec_lan_ac $ip"; done)
	EOF
    $IPT -I PREROUTING 1 ${IFNAME:+-i $IFNAME} -p tcp $EXT_ARGS $MATCH_SET \
        -m comment --comment "$TAG" -j SS_SPEC_WAN_AC
    if [ "$OUTPUT" = 1 ]; then
        $IPT -I OUTPUT 1 -p tcp $EXT_ARGS \
            -m comment --comment "$TAG" -j SS_SPEC_WAN_AC
    elif [ "$OUTPUT" = 2 ]; then
        ipset -! -R <<-EOF || return 1
			create ssr_gen_router hash:net hashsize 4096 maxelem 1000000
			$(gen_spec_iplist | sed -e "s/^/add ssr_gen_router /")
		EOF
        $IPT -N SS_SPEC_ROUTER &&
            $IPT -A SS_SPEC_ROUTER -m set --match-set ssr_gen_router dst -j RETURN &&
            $IPT -A SS_SPEC_ROUTER -j SS_SPEC_WAN_FW
        $IPT -I OUTPUT 1 -p tcp -m comment --comment "$TAG" -j SS_SPEC_ROUTER
    fi
    return $?
}

tp_rule() {
    [ -n "$TPROXY" ] || return 0
    ip rule add fwmark 0x01/0x01 table 100
    ip route add local 0.0.0.0/0 dev lo table 100
    local ipt="iptables -t mangle"
    $ipt -N SS_SPEC_TPROXY
    $ipt -A SS_SPEC_TPROXY -p udp --dport 53 -j RETURN
    $ipt -A SS_SPEC_TPROXY -p udp -d 0.0.0.0/8 -j RETURN
    $ipt -A SS_SPEC_TPROXY -p udp -d 10.0.0.0/8 -j RETURN
    $ipt -A SS_SPEC_TPROXY -p udp -d 127.0.0.0/8 -j RETURN
    $ipt -A SS_SPEC_TPROXY -p udp -d 169.254.0.0/16 -j RETURN
    $ipt -A SS_SPEC_TPROXY -p udp -d 172.16.0.0/12 -j RETURN
    $ipt -A SS_SPEC_TPROXY -p udp -d 192.168.0.0/16 -j RETURN
    $ipt -A SS_SPEC_TPROXY -p udp -d 224.0.0.0/4 -j RETURN
    $ipt -A SS_SPEC_TPROXY -p udp -d 240.0.0.0/4 -j RETURN
    $ipt -A SS_SPEC_TPROXY -p udp ! --dport 53 -d $server -j RETURN
    $ipt -A SS_SPEC_TPROXY -p udp $PROXY_PORTS -m set --match-set fplan src \
        -j TPROXY --on-port "$LOCAL_PORT" --tproxy-mark 0x01/0x01
    case "$RUNMODE" in
    router)
        $ipt -A SS_SPEC_TPROXY -p udp -m set --match-set gmlan src -m set ! --match-set china dst \
            -j TPROXY --on-port "$LOCAL_PORT" --tproxy-mark 0x01/0x01
        $ipt -A SS_SPEC_TPROXY -p udp $PROXY_PORTS -m set ! --match-set ss_spec_wan_ac dst \
            -j TPROXY --on-port "$LOCAL_PORT" --tproxy-mark 0x01/0x01
        ;;
    gfw)
        $ipt -A SS_SPEC_TPROXY -p udp -m set --match-set china dst -j RETURN
        $ipt -A SS_SPEC_TPROXY -p udp -m set --match-set gmlan src -m set ! --match-set china dst \
            -j TPROXY --on-port "$LOCAL_PORT" --tproxy-mark 0x01/0x01
        $ipt -A SS_SPEC_TPROXY -p udp -m set $PROXY_PORTS --match-set gfwlist dst \
            -j TPROXY --on-port "$LOCAL_PORT" --tproxy-mark 0x01/0x01
        ;;
    oversea)
        $ipt -A SS_SPEC_TPROXY -p udp $PROXY_PORTS -m set --match-set china dst \
            -j TPROXY --on-port "$LOCAL_PORT" --tproxy-mark 0x01/0x01
        ;;
    all)
        $ipt -A SS_SPEC_TPROXY -p udp $PROXY_PORTS -j TPROXY --on-port "$LOCAL_PORT" --tproxy-mark 0x01/0x01
        ;;
    esac
    $ipt -I PREROUTING 1 ${IFNAME:+-i $IFNAME} -p udp $EXT_ARGS $MATCH_SET \
        -m comment --comment "$TAG" -j SS_SPEC_TPROXY
    return $?
}

get_wan_ip() {
	cat <<-EOF | grep -E "^([0-9]{1,3}\.){3}[0-9]{1,3}"
		$server
		$SERVER
		$WAN_BP_IP
EOF
}

gen_iplist() {
	cat <<-EOF
		0.0.0.0/8
		10.0.0.0/8
		100.64.0.0/10
		127.0.0.0/8
		169.254.0.0/16
		172.16.0.0/12
		192.0.0.0/24
		192.0.2.0/24
		192.88.99.0/24
		192.168.0.0/16
		198.18.0.0/15
		198.51.100.0/24
		203.0.113.0/24
		224.0.0.0/4
		240.0.0.0/4
		255.255.255.255
		$(get_wan_ip)
		$(cat ${IGNORE_LIST:=/dev/null} 2>/dev/null)
EOF
}

gen_spec_iplist() {
    cat <<-EOF
		0.0.0.0/8
		10.0.0.0/8
		100.64.0.0/10
		127.0.0.0/8
		169.254.0.0/16
		172.16.0.0/12
		192.0.0.0/24
		192.0.2.0/24
		192.88.99.0/24
		192.168.0.0/16
		198.18.0.0/15
		198.51.100.0/24
		203.0.113.0/24
		224.0.0.0/4
		240.0.0.0/4
		255.255.255.255
		$(get_wan_ip)  
EOF
}

gen_include() {
    [ -n "$FWI" ] || return 0
    extract_rules() {
        echo "*$1"
        iptables-save -t $1 | grep SS_SPEC_ |
            sed -e "s/^-A \(OUTPUT\|PREROUTING\)/-I \1 1/"
        echo 'COMMIT'
    }
    cat <<-EOF >>$FWI
		iptables-save -c | grep -v "SS_SPEC" | iptables-restore -c
		iptables-restore -n <<-EOT
		$(extract_rules nat)
		$(extract_rules mangle)
		EOT
	EOF
    return 0
}

while getopts ":s:l:S:L:i:e:a:b:w:p:G:D:oOuUfgrczhd" arg; do
    case "$arg" in
    s)
        server=$OPTARG
        ;;
    l)
        local_port=$OPTARG
        ;;
    S)
        SERVER=$OPTARG
        ;;
    L)
        LOCAL_PORT=$OPTARG
        ;;
    i)
        IGNORE_LIST=$OPTARG
        ;;
    e)
        EXT_ARGS=$OPTARG
        ;;
    a)
        LAN_AC_IP=$OPTARG
        ;;
    b)
        WAN_BP_IP=$(for ip in $OPTARG; do echo $ip; done)
        ;;
    w)
        WAN_FW_IP=$OPTARG
        ;;
    p)
        LAN_FP_IP=$OPTARG
        ;;
    G)
        LAN_GM_IP=$OPTARG
        ;;
    D)
        PROXY_PORTS=$OPTARG
        ;;
    o)
        OUTPUT=1
        ;;
    O)
        OUTPUT=2
        ;;
    u)
        TPROXY=1
        ;;
    U)
        TPROXY=2
        ;;
    g)
        RUNMODE=gfw
        ;;
    r)
        RUNMODE=router
        ;;
    c)
        RUNMODE=oversea
        ;;
    z)
        RUNMODE=all
        ;;
    d)
        RUNMODE=direct
        ;;
    f)
        flush_r
        exit 0
        ;;
    h)
        usage 0
        ;;
    esac
done

if [ -z "$server" -o -z "$local_port" ]; then
    usage 2
fi

if [ "$TPROXY" == 1 ]; then
    SERVER=$server
    LOCAL_PORT=$local_port
elif [ "$TPROXY" == 2 ]; then
    : ${SERVER:?"You must assign an ip for the udp relay server."}
    : ${LOCAL_PORT:?"You must assign a port for the udp relay server."}
fi

flush_r && fw_rule && ipset_r && ac_rule && tp_rule && gen_include
[ "$?" == 0 ] || loger 3 "Start failed!"
exit $?
