#!/bin/bash

set -eu
# __SILENT=TRUE
. /opt/fox_utils/crab_sys.sh
[ "${1:---help}" = "--help" ] && sys::usage "$@"
sys::arg_parse "$@"
### --help Info: настройка фаервола для ВМ
### --help Info:
### --help Usage: vm_firewall <VM_NAME> <CMD>
### --help Usage:
### --help Usage: CMD:
### --help Usage:     start - создать цепочки правил фаервола для ВМ, создать правила
### --help Usage:     stop - удалить правила фаервола ВМ, удалить цепочки
### --help Usage:     migrate_out - правила фаервола для ВМ на время мигации на другую ноду
### --help Usage:     get_chain_name - вывести название цепочки ВМ, для хуков и прочих скриптов
### --help Usage:
### --help Example: vm_firewall myvm1 start
### --help Example: vm_firewall myvm1 stop
### --help Example: vm_firewall myvm1 do_get_chain_name
### --help Example:

declare VM_NAME VM_UUID VM_DNAT VM_IP CHAIN_NAME
declare CHAINS
declare CMD
declare COMMON_TCP_PORTS
declare COMMON_UDP_PORTS

OBJ="$1"
CMD="$2"
. /opt/fox_utils/fox_conf vm get "$OBJ"

echo TODO2 iptables -t nat -D PREROUTING -d  -p tcp --dport VM_SSH_PORT >&2
# -j DNAT --to-dest $VM_IP:${VM_SSH_PORT:-22} &>/dev/null || true
echo TODO2 !!! iptables -t nat -I PREROUTING -d -p tcp --dport SSH_PORT >&2
#  -j DNAT --to-dest $VM_IP:${VM_SSH_PORT:-22}

# в iptables -m multiport --dport максимум 15 портов за раз
# исключены порты 53,80,443
# для ftp проброс портов 10000:11999 TODO0 убрать после включения модулей ядра
# для upd исключены порты 20,21,22,25,110,143,993
# Порты для RDP: 5900:5910,6000:6001,3389,3350
# RadioMarvin icecast трансляция 8005,8006
COMMON_TCP_PORTS=(
"21,22,10000:11999,156,1433,1521" \
	"3306,5432,1500,2030,2031,2082,2083,2222,8083,8443" \
	"5672,6379,8000,8080,9300,9312,2087" \
	"5900:5910,6000:6001,3389,3350,8005,8006"
)
COMMON_UDP_PORTS=(
"156,5900:5910,6000:6001,3389,3350" \
	"1433,1521,3306,5432,1500,2030,2031,2082,2083,2222,8083,8443" \
	"5672,6379,8000,8080,9300,9312"
)

#######################################
# Сохранит список цепочек iptables в глобальной переменной
# Globals:
#   CHAINS - список цепочек iptables через пробел
#######################################
get_chains(){
	CHAINS="$(iptables-save)"
	CHAINS="$(grep ^:<<<"$CHAINS")"
	CHAINS="$(sed -E 's/:(\S+)\s.*/\1/g'<<<"$CHAINS")"
	return 0
}

#######################################
# Создаст цепочку iptables
# Globals:
#   CHAINS - список цепочек iptables через пробел
# Arguments:
#   --flush - отчистит цепочку, если она уже существует
#######################################
add_chain(){
	get_chains
	if ! grep -q "^$1$"<<<"$CHAINS"; then
		iptables -N "$1"
	fi
	[[ $@ =~ --flush ]] && iptables -F "$1"
	return 0
}

#######################################
# Удалит цепочку iptables
# Globals:
#   CHAINS - список цепочек iptables через пробел
#######################################
del_chain(){
	get_chains
	if grep -q "^$1$"<<<"$CHAINS"; then
		iptables -F "$1" || true
		iptables -X "$1" || true
	fi
	return 0
}

#######################################
# Удалит правило iptables, если оно существует
#######################################
del_rule(){
	set -e
	local rule check
	rule="$@"
	check="${rule/-I /-C }"
	check="${check/-A /-C }"
	check="${check/-D /-C }"
	if iptables $check; then iptables $rule; fi
	return 0
}

update_rule() {
	local rule="$*" del_rule
	del_rule="${rule/-I /-D }"
	del_rule="${del_rule/-A /-D }"

	iptables $del_rule &>/dev/null || true
	iptables $del_rule &>/dev/null || true
	echo "iptables $rule"
	iptables $rule
	return 0
}

update_vm_dnat_chain() {
	local rule proto src dport dst_port common_ports
	for rule in ${VM_DNAT:-}; do
		IFS=':' read -r proto src dport dst_port <<<"$rule"
		echo DNAT proto=$proto src=$src dport=$dport dst_port=$dst_port
		if [ "$proto" = "*" ]; then
			# TODO1 проверять, что данный админ ip не используется в других правилах/вм
			# dnat всего адреса, проброс только стандартных портов для протоколов tcp и udp
			for common_ports in "${COMMON_TCP_PORTS[@]}"; do
				update_rule -t nat -I "$CHAIN_NAME" -m addrtype --dst-type LOCAL \
					-s "$src" -p tcp -m multiport --dport "$common_ports" -j DNAT --to-dest "$VM_IP"
			done
			for common_ports in "${COMMON_UDP_PORTS[@]}"; do
				update_rule -t nat -I "$CHAIN_NAME" -m addrtype --dst-type LOCAL \
					-s "$src" -p udp -m multiport --dport "$common_ports" -j DNAT --to-dest "$VM_IP"
			done
		else
			# dnat порта
			update_rule -t nat -I "$CHAIN_NAME" -m addrtype --dst-type LOCAL \
				-s $src -p $proto --dport $dport \
				-j DNAT --to-dest $VM_IP:$dst_port
		fi
	done
	return 0
}

update_output_drop(){
	# на время стопа мы дропали пакеты чтоб не показывать nginx 500 error
	local cmd="$1"
	local ip
	local ipt_rules
	for ip in $VM_IP; do
		if [ "$cmd" = start_drop ]; then
			echo "vm_firewall: Включаем DROP в OUTPUT для $ip"
			set -o pipefail
			ipt_rules=$(iptables -nvL OUTPUT | grep "DROP" || true)
			set +o pipefail
			if [[ "$ipt_rules" == *" $ip "* ]]; then
				echo "vm_firewall: DROP уже был установлен!"
				continue
			fi
			iptables -I OUTPUT -d "$ip" -m state --state NEW -j DROP
		elif [ "$cmd" = stop_drop ]; then
			echo "vm_firewall: Выключаем DROP в OUTPUT для $ip"
			set -o pipefail
			ipt_rules=$(iptables -nvL OUTPUT | grep "DROP" || true)
			set +o pipefail
			if [[ "$ipt_rules" != *" $ip "* ]]; then
				echo "vm_firewall: DROP не был установлен!"
				continue
			fi
			iptables -D OUTPUT -d "$ip" -m state --state NEW -j DROP
		else
			echo "Error! Неизвестная команда update_output_drop $cmd"
		fi
	done
	return 0
}

update_dnat_masquerade(){
	# см. /проект_cloud_и_vps_2.0:в_работе:2021-12-01_dnat_и_masquerade_при_миграции
	# При миграции vm на другую ноду мы оставляем правила DNAT
	# ssh_dnat, admin_ip и тп, так как клиент может какое-то время еще
	# ходить на старую ноду (dns кеш)
	# Чтобы эти DNAT работали, мы делаем MASQUERADE
	# Он временный (TODO: сделать чистку, мб переделать на ipset)
	# В таком варианте у клиента при заходе на впс будет ip-адрес ноды
	# но считаем, что это лучше, чем проблемы с подключением
	local cmd="$1"
	local ip
	for ip in $VM_IP; do
		if [ "$cmd" = start_dnat_masquerade ]; then
			echo "vm_firewall: Включаем MASQUERADE для $ip"
			# Обязательно -A, т.к. сначала в цепочке идут исключающие правила
			update_rule -t nat -A MIGRATE_OUT -d "$VM_IP" -j MASQUERADE
		elif [ "$cmd" = stop_dnat_masquerade ]; then
			echo "vm_firewall: Выключаем MASQUERADE для $ip"
			iptables -t nat -D MIGRATE_OUT -d "$VM_IP" -j MASQUERADE &>/dev/null || true
			iptables -t nat -D MIGRATE_OUT -d "$VM_IP" -j MASQUERADE &>/dev/null || true
		else
			echo "Error! Неизвестная команда update_dnat_masquerade $cmd"
		fi
	done
	return 0
}

do_start() {
	## FORWARD
	add_chain "$CHAIN_NAME" --flush
	iptables -N "$CHAIN_NAME" &>/dev/null || true
	for ip in $VM_IP; do
		# Ограничиваем адресами ВМ, чтобы -A $CHAIN_NAME -j DROP не перекрыл весь FORWARD ноды
		update_rule -I VM_FORWARD -s "$VM_IP" -j "$CHAIN_NAME"
		update_rule -I VM_FORWARD -d "$VM_IP" -j "$CHAIN_NAME"
	done
	iptables -F "$CHAIN_NAME"

	## NAT
	iptables -t nat -N "$CHAIN_NAME" &>/dev/null || true
	update_rule -t nat -I VM_DNAT -j "$CHAIN_NAME"
	iptables -t nat -F "$CHAIN_NAME"
	# Если вернулись после миграции на страрую ноду и еще не удалили правила
	update_dnat_masquerade "stop_dnat_masquerade"
	update_vm_dnat_chain
	conntrack -D -g $VM_IP || true

	update_output_drop "stop_drop"
	return 0
}

do_stop() {
	## FORWARD
	for ip in $VM_IP; do
		del_rule -D VM_FORWARD -s "$VM_IP" -j "$CHAIN_NAME"
		del_rule -D VM_FORWARD -d "$VM_IP" -j "$CHAIN_NAME"
	done
	del_chain "$CHAIN_NAME"

	## NAT
	iptables -t nat -F "$CHAIN_NAME" &>/dev/null || true
	iptables -t nat -D VM_DNAT -j "$CHAIN_NAME" || true
	iptables -t nat -D VM_DNAT -j "$CHAIN_NAME" &>/dev/null || true
	iptables -t nat -X "$CHAIN_NAME" || true

	update_output_drop "start_drop"
	return 0
}

do_migrate_out() {
	do_stop
	# Правила VM_DNAT запишем в цепочку MIGRATE_DNAT, чтобы потом очистить её по крону или вручную
	CHAIN_NAME=MIGRATE_DNAT
	update_vm_dnat_chain
	update_dnat_masquerade "start_dnat_masquerade"
	return 0
}


#######################################
# Вернёт имя цепочки для ВМ
#
# Имя цепочки можно использовать в start/stop хуках ВМ
# или в сторонних скриптах
#######################################
do_get_chain_name() {
	echo "$CHAIN_NAME"
	return 0
}


main() {
	/opt/fox_utils/crab_lock "/tmp/vm_firewall.lock" $$
	# в iptables название цепочки максимум 28 символов
	CHAIN_NAME="vm_${VM_NAME:0:16}_${VM_UUID:0:8}"
	"do_$CMD" "$@"
	/opt/fox_utils/crab_unlock "/tmp/vm_firewall.lock" $$
	return 0
}

main "$@"
exit 0