#!/bin/bash
set -euo pipefail
[[ -n "${DEBUG:-}" ]] && set -x

LANG=en_US.UTF-8
LANGUAGE=en_US:
LC_ALL=en_US.UTF-8
PATH="/bin:/usr/bin:/sbin:/usr/sbin"

GREP_REGEXP_INSTANCE_NAME="[-_.[:alnum:]]\\+"
DEFAULT_PROTO=tcp

ufw_docker_agent=ufw-docker-agent
ufw_docker_agent_image="${UFW_DOCKER_AGENT_IMAGE:-chaifeng/${ufw_docker_agent}:200812}"

function ufw-docker--status() {
    ufw-docker--list "$GREP_REGEXP_INSTANCE_NAME"
}

function ufw-docker--list() {
    local INSTANCE_NAME="$1"
    local INSTANCE_PORT="${2:-}"
    local PROTO="${3:-${DEFAULT_PROTO}}"

    if [[ -z "$INSTANCE_PORT" ]]; then
        INSTANCE_PORT="[[:digit:]]\\+"
        PROTO="\\(tcp\\|udp\\)"
    fi
    ufw status numbered | grep "# allow ${INSTANCE_NAME}\\( ${INSTANCE_PORT}\\/${PROTO}\\)\\?\$"
}

function ufw-docker--list-number() {
    ufw-docker--list "$@" | sed -e 's/^\[[[:blank:]]*\([[:digit:]]\+\)\].*/\1/'
}

function ufw-docker--delete() {
    for UFW_NUMBER in $(ufw-docker--list-number "$@" | sort -rn); do
        echo "delete \"$UFW_NUMBER\""
        echo y | ufw delete "$UFW_NUMBER" || true
    done
}

function ufw-docker--allow() {
    local INSTANCE_NAME="$1"
    local INSTANCE_PORT="$2"
    local PROTO="$3"

    docker inspect "$INSTANCE_NAME" &>/dev/null ||
        die "Docker instance \"$INSTANCE_NAME\" doesn't exist."

    mapfile -t INSTANCE_IP_ADDRESSES < <(docker inspect --format='{{range .NetworkSettings.Networks}}{{.IPAddress}}{{"\n"}}{{end}}' "$INSTANCE_NAME" 2>/dev/null | remove_blank_lines)

    [[ -z "${INSTANCE_IP_ADDRESSES:-}" ]] && die "Could not find a running instance \"$INSTANCE_NAME\"."

    mapfile -t PORT_PROTO_LIST < <(docker inspect --format='{{range $p, $conf := .NetworkSettings.Ports}}{{with $conf}}{{$p}}{{"\n"}}{{end}}{{end}}' "$INSTANCE_NAME" | remove_blank_lines)

    if [[ -z "${PORT_PROTO_LIST:-}" ]]; then
        err "\"$INSTANCE_NAME\" doesn't have any published ports."
        return 1
    fi

    RETVAL=1
    for PORT_PROTO in "${PORT_PROTO_LIST[@]}"; do
        if [[ -z "$INSTANCE_PORT" || "$PORT_PROTO" = "${INSTANCE_PORT}/${PROTO}" ]]; then
            for IP in "${INSTANCE_IP_ADDRESSES[@]}"; do
                ufw-docker--add-rule "$INSTANCE_NAME" "$IP" "${PORT_PROTO%/*}" "${PORT_PROTO#*/}"
                RETVAL="$?"
            done
        fi
    done
    if [[ "$RETVAL" -ne 0 ]]; then
        err "Fail to add rule(s), cannot find the published port ${INSTANCE_PORT}/${PROTO} of instance \"${INSTANCE_NAME}\" or cannot update outdated rule(s)."
    fi
    return "$RETVAL"
}

function ufw-docker--add-service-rule() {
    declare service_id="$1"
    declare port="${2%/*}"
    declare proto="${2#*/}"

    declare target_ip_port
    target_ip_port="$(iptables -t nat -L DOCKER-INGRESS | grep -E "^DNAT\\s+${proto}\\s+.+\\sto:[.0-9]+:${port}\$" | grep -Eo "[.0-9]+:${port}\$")"

    [[ -z "$target_ip_port" ]] && die "Could not find VIP of service ${service_id}."

    ufw-docker--add-rule "$service_id" "${target_ip_port%:*}" "$port" "$proto"
}

function ufw-docker--add-rule() {
    local INSTANCE_NAME="$1"
    local INSTANCE_IP_ADDRESS="$2"
    local PORT="$3"
    local PROTO="$4"

    declare comment

    echo "allow ${INSTANCE_NAME} ${PORT}/${PROTO}"
    typeset -a UFW_OPTS
    UFW_OPTS=(route allow proto "${PROTO}"
              from any to "$INSTANCE_IP_ADDRESS")
    comment="allow ${INSTANCE_NAME}"
    [[ -n "$PORT" ]] && {
        UFW_OPTS+=(port "${PORT}")
        comment="$comment ${PORT}/${PROTO}"
    }
    UFW_OPTS+=(comment "$comment")

    if ufw-docker--list "$INSTANCE_NAME" "$PORT" "$PROTO" &>/dev/null; then
        ufw --dry-run "${UFW_OPTS[@]}" | grep "^Skipping" && return 0
        err "Remove outdated rule."
        ufw-docker--delete "$INSTANCE_NAME" "$PORT" "$PROTO"
    fi
    echo ufw "${UFW_OPTS[@]}"
    ufw "${UFW_OPTS[@]}"
}

function ufw-docker--instance-name() {
    local INSTANCE_ID="$1"
    {
        {
            docker inspect --format='{{.Name}}' "$INSTANCE_ID" 2>/dev/null | sed -e 's,^/,,' |
                grep "^${GREP_REGEXP_INSTANCE_NAME}\$" 2>/dev/null
        } || echo -n "$INSTANCE_ID";
    } | remove_blank_lines
}

function ufw-docker--service() {
    declare service_action="${1:-help}"
    case "$service_action" in
        delete)
            shift || true
            if [[ "${1:?Invalid 'delete' command syntax.}" != "allow" ]]; then
                die "\"delete\" command only support removing allowed rules"
            fi
            shift || true
            declare service_id_or_name="${1:?Missing swarm service name or service ID}"

            "ufw-docker--service-${service_action}" "${service_id_or_name}"
            ;;
        allow)
            shift || true
            declare service_id_or_name="${1:?Missing swarm service name or service ID}"
            declare service_port="${2:?Missing the port number, such as '80/tcp'.}"

            "ufw-docker--service-${service_action}" "${service_id_or_name}" "${service_port}"
            ;;
        *)
            ufw-docker--help
            ;;
    esac
}

function ufw-docker--get-service-id() {
    declare service_name="$1"
    docker service inspect "${service_name}" --format "{{.ID}}"
}

function ufw-docker--get-service-name() {
    declare service_name="$1"
    docker service inspect "${service_name}" --format "{{.Spec.Name}}"
}

function ufw-docker--service-allow() {
    declare service_name="$1"
    declare service_port="$2"
    declare service_proto=tcp

    if [[ -n "$service_port" ]] &&
           ! grep -E '^[0-9]+(/(tcp|udp))?$' <<< "$service_port" &>/dev/null; then
        die "Invalid port syntax: $service_port"
        return 1
    fi

    if [[ "$service_port" = */* ]]; then
        service_proto="${service_port#*/}"
        service_port="${service_port%/*}"
    fi

    declare service_id
    service_id="$(ufw-docker--get-service-id "${service_name}")"
    [[ -z "${service_id:-}" ]] && die "Could not find service \"$service_name\""

    service_name="$(ufw-docker--get-service-name "${service_name}")"

    exec 9< <(docker service inspect "$service_name" \
                     --format '{{range .Endpoint.Spec.Ports}}{{.PublishedPort}} {{.TargetPort}}/{{.Protocol}}{{"\n"}}{{end}}')
    while read -u 9 -r port target_port; do
        if [[ "$target_port" = "${service_port}/${service_proto}" ]]; then
            declare service_env="ufw_public_${service_id}=${service_name}/${port}/${service_proto}"
            break;
        fi
    done
    exec 9<&-

    [[ -z "${service_env:-}" ]] && die "Service $service_name does not publish port $service_port."

    if ! docker service inspect "$ufw_docker_agent" &>/dev/null; then
        err "Not found ufw-docker-agent service, creating ..."
        docker service create --name "$ufw_docker_agent" --mode global \
               --mount type=bind,source=/var/run/docker.sock,target=/var/run/docker.sock \
               --mount type=bind,source=/etc/ufw,target=/etc/ufw,readonly=true \
               --env ufw_docker_agent_image="${ufw_docker_agent_image}" \
               --env DEBUG="${DEBUG:-}" \
               --env "${service_env}" \
               "${ufw_docker_agent_image}"
    else
        declare -a service_env_list
        service_env_list+=(--env-add "${service_env}")

        exec 8< <(ufw-docker--get-env-list)
        while read -u 8 -r id value; do
            [[ "$id" = "$service_id" ]] && continue
            [[ "$value" = "${service_name}"/* ]] && service_env_list+=(--env-rm "ufw_public_${id}")
        done
        exec 8<&-

        docker service update --update-parallelism=0 \
               --env-add ufw_docker_agent_image="${ufw_docker_agent_image}" \
               --env-add DEBUG="${DEBUG:-}" \
               "${service_env_list[@]}" \
               --image "${ufw_docker_agent_image}" \
               "${ufw_docker_agent}"
    fi
}

function ufw-docker--get-env-list() {
    docker service inspect "${ufw_docker_agent}" \
           --format '{{range $k,$v := .Spec.TaskTemplate.ContainerSpec.Env}}{{ $v }}{{"\n"}}{{end}}' |
        sed -e '/^ufw_public_/!d' \
            -e 's/^ufw_public_//' \
            -e 's/=/ /'
}

function ufw-docker--service-delete() {
    declare service_name="$1"

    exec 8< <(ufw-docker--get-env-list)
    while read -u 8 -r id value; do
        if [[ "$id" = "$service_name" ]] || [[ "$value" = "${service_name}"/* ]]; then
            declare service_id="$id"
            service_name="${value%%/*}"
            declare service_env="ufw_public_${service_id}=${service_name}/deny"
            break;
        fi
    done
    exec 8<&-

    [[ -z "${service_env:-}" ]] && die "Could not find service \"$service_name\""

    docker service update --update-parallelism=0 \
           --env-add ufw_docker_agent_image="${ufw_docker_agent_image}" \
           --env-add "${service_env}" \
           --image "${ufw_docker_agent_image}" \
           "${ufw_docker_agent}"
}

function ufw-docker--raw-command() {
    ufw "$@"
}

after_rules="/etc/ufw/after.rules"

function ufw-docker--check() {
  err "\\n########## iptables -n -L DOCKER-USER ##########"
  iptables -n -L DOCKER-USER

  err "\\n\\n########## diff $after_rules ##########"
  ufw-docker--check-install && err "\\nCheck done."
}

declare -a files_to_be_deleted

function rm-on-exit() {
  [[ $# -gt 0 ]] && files_to_be_deleted+=("$@")
}

function on-exit() {
  for file in "${files_to_be_deleted[@]:-}"; do
    [[ -f "$file" ]] && rm -r "$file"
  done
  files_to_be_deleted=()
}

trap on-exit EXIT INT TERM QUIT ABRT ERR

function ufw-docker--check-install() {
	after_rules_tmp="${after_rules_tmp:-$(mktemp)}"
	rm-on-exit "$after_rules_tmp"

	sed "/^# BEGIN UFW AND DOCKER/,/^# END UFW AND DOCKER/d" "$after_rules" > "$after_rules_tmp"
	>> "${after_rules_tmp}" cat <<-\EOF
	# BEGIN UFW AND DOCKER
	*filter
	:ufw-user-forward - [0:0]
	:ufw-docker-logging-deny - [0:0]
	:DOCKER-USER - [0:0]
	-A DOCKER-USER -j ufw-user-forward

	-A DOCKER-USER -j RETURN -s 10.0.0.0/8
	-A DOCKER-USER -j RETURN -s 172.16.0.0/12
	-A DOCKER-USER -j RETURN -s 192.168.0.0/16

	-A DOCKER-USER -p udp -m udp --sport 53 --dport 1024:65535 -j RETURN

	-A DOCKER-USER -j ufw-docker-logging-deny -p tcp -m tcp --tcp-flags FIN,SYN,RST,ACK SYN -d 192.168.0.0/16
	-A DOCKER-USER -j ufw-docker-logging-deny -p tcp -m tcp --tcp-flags FIN,SYN,RST,ACK SYN -d 10.0.0.0/8
	-A DOCKER-USER -j ufw-docker-logging-deny -p tcp -m tcp --tcp-flags FIN,SYN,RST,ACK SYN -d 172.16.0.0/12
	-A DOCKER-USER -j ufw-docker-logging-deny -p udp -m udp --dport 0:32767 -d 192.168.0.0/16
	-A DOCKER-USER -j ufw-docker-logging-deny -p udp -m udp --dport 0:32767 -d 10.0.0.0/8
	-A DOCKER-USER -j ufw-docker-logging-deny -p udp -m udp --dport 0:32767 -d 172.16.0.0/12

	-A DOCKER-USER -j RETURN

	-A ufw-docker-logging-deny -m limit --limit 3/min --limit-burst 10 -j LOG --log-prefix "[UFW DOCKER BLOCK] "
	-A ufw-docker-logging-deny -j DROP

	COMMIT
	# END UFW AND DOCKER
	EOF

	diff -u --color=auto "$after_rules" "$after_rules_tmp"
}

function ufw-docker--install() {
  if ! ufw-docker--check-install; then
	  local after_rules_bak
    after_rules_bak="${after_rules}-ufw-docker~$(date '+%Y-%m-%d-%H%M%S')~"
	  err "\\nBacking up $after_rules to $after_rules_bak"
	  cp "$after_rules" "$after_rules_bak"
	  cat "$after_rules_tmp" > "$after_rules"
    err "Please restart UFW service manually by using the following command:"
    if type systemctl &>/dev/null; then
      err "    sudo systemctl restart ufw"
    else
      err "    sudo service ufw restart"
    fi
  fi
}

function ufw-docker--help() {
	cat <<-EOF >&2
	Usage:
	  ufw-docker <list|allow> [docker-instance-id-or-name [port[/tcp|/udp]]]
	  ufw-docker delete allow [docker-instance-id-or-name [port[/tcp|/udp]]]

	  ufw-docker service allow <swarm-service-id-or-name <port</tcp|/udp>>>
	  ufw-docker service delete allow <swarm-service-id-or-name>

	  ufw-docker <status|install|check|help>

	Examples:
	  ufw-docker help

	  ufw-docker check           # Check the installation of firewall rules
	  ufw-docker install         # Install firewall rules

	  ufw-docker status

	  ufw-docker list httpd


	  ufw-docker allow httpd
	  ufw-docker allow httpd 80
	  ufw-docker allow httpd 80/tcp

	  ufw-docker delete allow httpd
	  ufw-docker delete allow httpd 80/tcp


	  ufw-docker service allow httpd 80/tcp

	  ufw-docker service delete allow httpd
	EOF
}

function remove_blank_lines() {
    sed '/^[[:blank:]]*$/d'
}

function err() {
    echo -e "$@" >&2
}

function die() {
    err "ERROR:" "$@"
    exit 1
}

# __main__

if ! ufw status 2>/dev/null | grep -Fq "Status: active" ; then
    die "UFW is disabled or you are not root user."
fi

ufw_action="${1:-help}"

case "$ufw_action" in
    delete)
        shift || true
        if [[ "${1:?Invalid 'delete' command syntax.}" != "allow" ]]; then
            die "\"delete\" command only support removing allowed rules"
        fi
        ;&
    list|allow)
        shift || true

        INSTANCE_ID="${1:?Docker instance name/ID cannot be empty.}"
        INSTANCE_NAME="$(ufw-docker--instance-name "$INSTANCE_ID")"
        shift || true

        INSTANCE_PORT="${1:-}"
        if [[ -n "$INSTANCE_PORT" && ! "$INSTANCE_PORT" =~ [0-9]+(/(tcp|udp))? ]]; then
            die "invalid port syntax: \"$INSTANCE_PORT\"."
        fi

        PROTO="$DEFAULT_PROTO"
        if [[ "$INSTANCE_PORT" = */udp ]]; then
            PROTO=udp
        fi

        INSTANCE_PORT="${INSTANCE_PORT%/*}"

        "ufw-docker--$ufw_action" "$INSTANCE_NAME" "$INSTANCE_PORT" "$PROTO"
        ;;
    service|raw-command|add-service-rule)
        shift || true
        "ufw-docker--$ufw_action" "$@"
        ;;
    status|install|check)
        ufw-docker--"$ufw_action"
        ;;
    *)
        ufw-docker--help
        ;;
esac