#!/bin/bash -e

HOSTGROUP=disk_db_misc_01
PORT=27018

fail() {
    echo -e "\n$1 \e[1m\e[31mFail\e[0m"
    exit 1
}

success() {
    echo -e "\n$1 \e[1m\e[32mSuccess\e[0m"
    exit 0
}

execute() {
    $1 || fail "$1"
    success "$1"
}

hide() {
  # create ipset list if it does not exist
  ipset -! create allowed_hosts hash:ip family inet6

  # get hosts from conductor group and resolve them 
  for h in $(curl -s https://c.yandex-team.ru/api/groups2hosts/$HOSTGROUP):
    do
      host_addr=$(dig -t AAAA +short "$h")
      if [[ -n $host_addr ]]; then
        # add resolved host ipv6 address to ipset list
        ipset -! add allowed_hosts "$host_addr"
        not_empty=True # ipset list is not empty now
      fi
    done

  if [[ -n $not_empty ]]; then
    # insert rule if it does not already exist
    ip6tables -C MANUAL -p tcp --dport $PORT -m set --match-set allowed_hosts src -j ACCEPT >/dev/null 2>&1 || ip6tables -A MANUAL -p tcp --dport $PORT -m set --match-set allowed_hosts src -j ACCEPT
    ip6tables -C MANUAL -p tcp --dport $PORT -j REJECT --reject-with tcp-reset >/dev/null 2>&1 || ip6tables -A MANUAL -p tcp --dport $PORT -j REJECT --reject-with tcp-reset
    # list hosts from  ipset and iptables rules
    ipset -r list allowed_hosts
    ip6tables -vnL MANUAL
  else
    echo "Empty hostlist after resolving conductor group"
    return 1
  fi
}

unhide() {
  # check that iptables rule exists and remove it
  ip6tables -C MANUAL -p tcp --dport $PORT -j REJECT --reject-with tcp-reset >/dev/null 2>&1 && ip6tables -D MANUAL -p tcp --dport $PORT -j REJECT --reject-with tcp-reset
  ip6tables -C MANUAL -p tcp --dport $PORT -m set --match-set allowed_hosts src -j ACCEPT >/dev/null 2>&1 && ip6tables -D MANUAL -p tcp --dport $PORT -m set --match-set allowed_hosts src -j ACCEPT
  # check that ipset list exists and remove it
  ipset list allowed_hosts >/dev/null 2>&1 && ipset destroy allowed_hosts
  # show contents of iptabales MANUAL chain
  ip6tables -vnL MANUAL
}

if [[ "$1" == "unhide" ]]; then
  execute unhide 
elif [[ "$1" == "hide" ]]; then
  execute hide
else
  echo "Usage: $0 <hide|unhide>"
  exit 1
fi

