###############################################################
# @description:                                               #
#   Add Custom rule to UFW for limiting ssh(4242/tcp)         #
#                                                             #
# @author: Bram Prieshof                                      #
###############################################################

#Sed in a function to detect line from $FindLine and insert the content of $AddLine on a new line above the match
UpdateFile () {
    file="$1"
    sed -i 's/'"$FindLine"'/'"$AddLine"'\n&/g' $file
    unset FindLine AddLine
}

#Delete Existng SSH(4242/tcp) rule(s)
ufw delete limit 4242/tcp

#Add needed filters (IPv4)
FindLine="# End required lines"
AddLine=":ufw-user-limit - [0:0]\n:ufw-user-limit-accept - [0:0]"
UpdateFile /etc/ufw/after.rules

#Add needed filters (IPv6)
FindLine="# End required lines"
AddLine=":ufw6-user-limit - [0:0]\n:ufw6-user-limit-accept - [0:0]"
UpdateFile /etc/ufw/after6.rules

#Add custom SSH(4242/tcp) limit rule (IPv4)
FindLine="# don't delete the 'COMMIT' line or these rules won't be processed"
AddLine="### SSH limit tcp\n-A ufw-after-input -p tcp --dport 4242 -m conntrack --ctstate NEW -m recent --set\n-A ufw-after-input -p tcp --dport 4242 -m conntrack --ctstate NEW -m recent --update --seconds 30 --hitcount 15 -j ufw-user-limit\n-A ufw-after-input -p tcp --dport 4242 -j ufw-user-limit-accept\n"
UpdateFile /etc/ufw/after.rules

#Add custom SSH(4242/tcp) limit rule (IPv6)
FindLine="# don't delete the 'COMMIT' line or these rules won't be processed"
AddLine="### SSH limit tcp\n-A ufw6-after-input -p tcp --dport 4242 -m conntrack --ctstate NEW -m recent --set\n-A ufw6-after-input -p tcp --dport 4242 -m conntrack --ctstate NEW -m recent --update --seconds 30 --hitcount 15 -j ufw6-user-limit\n-A ufw6-after-input -p tcp --dport 4242 -j ufw6-user-limit-accept\n"
UpdateFile /etc/ufw/after6.rules

#Reload ufw rules
ufw reload