#!/usr/bin/bash

function die() {
    echo -e "Error: $@"
    exit 1
}

function get_option() {
    args="$1"
    option="$2"

    echo $(echo $args | grep -o -P "(?<=--$option |-$option ).*?(?:(?= -| --)|$)")
}

function check_chain() {
    protocol="$1"
    chain="sodalite-$2"

    function check_passthrough() {
        protocol="$1"
        args="$2"

        if [[ $(firewall-cmd --permanent --direct --query-passthrough $protocol $args) == "no" ]]; then
            # TODO: Remove passthrough if not needed
            firewall-cmd --permanent --direct --add-passthrough $protocol $args
        fi
    }

    if [[ $(firewall-cmd --permanent --direct --query-chain $protocol filter $chain) == "no" ]]; then
        # TODO: Remove chain if not needed
        firewall-cmd --permanent --direct --add-chain $protocol filter $chain
    fi
    
    case ${chain} in
        "sodalite-allow") check_passthrough $protocol "--append sodalite-allow --jump ACCEPT" ;;
        "sodalite-deny") check_passthrough $protocol "--append sodalite-deny --jump DROP" ;;
        "sodalite-reject") check_passthrough $protocol "--append sodalite-reject --jump REJECT" ;;
        "sodalite-limit")
            check_passthrough $protocol "--append sodalite-limit --match limit --limit 3/min --limit-burst 10 --jump ACCEPT"
            check_passthrough $protocol "--append sodalite-limit --jump DROP"
            ;;
        *)
            firewall-cmd --permanent --direct --remove-chain $protocol filter $chain
    esac
}

function get_status() {
    fwc_state="$(firewall-cmd --state)"

    case ${fwc_state} in
        "running") echo "active" ;;
        *) echo "inactive" ;;
    esac
}

function list_rules() {
    echo -e "Status: $(get_status)"
    echo -e ""
    echo -e "     To                         Action      From"
    echo -e "     --                         ------      ----"

    id=0
    firewall-cmd --permanent --direct --get-all-rules | while read line; do
        id=$((id+1))
    
        if [[ $line =~ ((ipv4|ipv6) filter (INPUT|OUTPUT) ([[:digit:]]+) (.*)) ]]; then
            args="${BASH_REMATCH[5]}"
            direction="${BASH_REMATCH[3]}"
            protocol="${BASH_REMATCH[2]}"

            ipt_jump="$(get_option "$args" "j")"
            ipt_port="$(get_option "$args" "dport")"
            ipt_protocol="$(get_option "$args" "p")"
            ipt_source="$(get_option "$args" "s")"

            if [[ $ipt_jump == "sodalite"* ]]; then
                ufw_action=""
                ufw_direction=""
                ufw_from="Anywhere"
                ufw_id=$id
                ufw_to="$ipt_port/$ipt_protocol"

                case "$ipt_jump" in
                    "sodalite-allow") ufw_action="ALLOW" ;;
                    "sodalite-deny") ufw_action="DENY" ;;
                    "sodalite-limit") ufw_action="LIMIT" ;;
                    "sodalite-reject") ufw_action="REJECT" ;;
                esac

                case "$direction" in
                    "INPUT") ufw_direction="IN" ;;
                    "OUTPUT") ufw_direction="OUT" ;;
                esac

                if [[ -n $ipt_source ]]; then
                    ufw_from="$ipt_source"
                fi

                if [[ $protocol == "ipv6" ]]; then
                    ufw_to+=" (v6)"
                    ufw_from+=" (v6)"
                fi

                echo -e "[$ufw_id] $ufw_to  $ufw_action $ufw_direction  $ufw_from  +$line"
            fi
        fi
    done
}

function add_rule() {
    ufw_options="$@"

    if [[ "$ufw_options" =~ ((allow|deny|reject|limit) (in|out) proto (tcp|udp) to ([0-9\.\/]{9,18}|[\:\/0-9A-Fa-f]{4,43}|any) port ([0-9]{1,5}) from ([0-9\.\/]{9,18}|[\:\/0-9A-Fa-f]{4,43}|any)) ]]; then
        ufw_action="${BASH_REMATCH[2]}"
        ufw_direction="${BASH_REMATCH[3]}"
        ufw_from="${BASH_REMATCH[7]}"
        ufw_port="${BASH_REMATCH[6]}"
        ufw_protocol="${BASH_REMATCH[4]}"
        ufw_to="${BASH_REMATCH[5]}"
        
        fwd_args="-p $ufw_protocol --dport $ufw_port -j sodalite-$ufw_action"
        fwd_direction=""
        
        case ${ufw_direction} in
            "in") fwd_direction="INPUT" ;;
            "out") fwd_direction="OUTPUT" ;;
        esac
        
        if { [[ $ufw_from != "0.0.0.0/0" ]] && [[ $ufw_from != "::/0" ]] && [[ $ufw_from != "any" ]]; }; then
            fwd_args+=" -s $ufw_from"
        fi

        if [[ $ufw_to =~ ([0-9\.\/]{9,18}|any) ]]; then
            check_chain ipv4 $ufw_action
            firewall-cmd --permanent --direct --add-rule ipv4 filter $fwd_direction 0 $fwd_args || die "Failed to add rule ($fwd_args)"
        fi

        if [[ $ufw_to =~ ([\:\/0-9A-Fa-f]{4,43}|any) ]]; then
            check_chain ipv6 $ufw_action
            firewall-cmd --permanent --direct --add-rule ipv6 filter $fwd_direction 0 $fwd_args || die "Failed to add rule ($fwd_args)"
        fi

        firewall-cmd --reload || die "Failed to reload"
    fi
}

function del_rule() {
    id=$1

    list_rules | while read line; do
        if [[ $line =~ (\[([[:digit:]]+)\] [^*]*\+(.*)) ]]; then
            ufw_id="${BASH_REMATCH[2]}"
            ufw_args="${BASH_REMATCH[3]}"

            if [[ $ufw_id == $id ]]; then
                firewall-cmd --permanent --direct --remove-rule $ufw_args
                firewall-cmd --reload
                break
            fi
        fi
    done
}

fw_daemon="firewalld"

if [[ -x "$(command -v ufw)" ]]; then
    fw_daemon="ufw"
fi

case ${fw_daemon} in
    "firewalld")
        while getopts "12345:6:" OPTION; do
            case ${OPTION} in
                2)
                    systemctl enable --now firewalld
                    ;;
                3)
                    systemctl disable --now firewalld
                    ;;
                4|1)
                    list_rules
                    ;;
                5)
                    add_rule $2
                    ;;
                6)
                    del_rule $OPTARG
                    ;;
                *)
                    die "Not implemented"
                    ;;
            esac
        done
        ;;
    "ufw")
        # https://github.com/elementary/switchboard-plug-security-privacy/blob/1119b7fd1441f1b06d37caea19f38216ad1f24da/data/security-privacy-plug-helper#L3-L24
        while getopts "12345:6:" OPTION; do
            case ${OPTION} in
                1)
                    LANGUAGE=C
                    export LANGUAGE
                    ufw status;;
                2)
                    ufw --force enable;;
                3)
                    ufw disable;;
                4)
                    LANGUAGE=C
                    export LANGUAGE
                    ufw status numbered;;
                5)
                    ufw $2;;
                6)
                    ufw --force delete $OPTARG;;
                \?)
                    exit 1;;
            esac
        done
        ;;
esac
