From 9cc63a9c22ae5c0ecac1dafe86f2e6610d3e784f Mon Sep 17 00:00:00 2001
From: AR2000 <ar2000@ar2000.me>
Date: Wed, 18 Dec 2024 22:31:50 +0100
Subject: [PATCH] more env var

---
 README.md  |  1 +
 bouncer.sh | 30 +++++++++++++-----------------
 2 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index e96576d..2672db6 100644
--- a/README.md
+++ b/README.md
@@ -10,6 +10,7 @@ A Crowdsec firewall bouncer based on [crowdsec-custom-bouncer](https://github.co
 - `IPTABLES_COMMENT` : if set, iptables will load the `comment` module and use the reason argument passed by [crowdsec-custom-bouncer](https://github.com/crowdsecurity/cs-custom-bouncer) as the comment
 - `IPTABLES_INSERT` : insert new firewall rules at a fixed position
 - `API_KEY` : crowdsec api key
+- `IPTABLES_USE_CHAIN` : chain to use
 
 ## Configuration :
 Set your custom configuration in the `crowdsec-custom-bouncer.yaml.local` file and mount it as a volume at the root of the container. See [Overriding values](https://docs.crowdsec.net/docs/next/configuration/crowdsec_configuration/#overriding-values) or [bouncers/custom](https://docs.crowdsec.net/u/bouncers/custom) for information about the configuration.
diff --git a/bouncer.sh b/bouncer.sh
index 151342f..54c05de 100755
--- a/bouncer.sh
+++ b/bouncer.sh
@@ -2,54 +2,50 @@
 #
 # Script to add /remove IPs to iptables
 
-[[ -z "${IPTABLES_COMMENT}" ]] && commment="-m comment --comment \"$4\"" || comment=""
+[[ -n "${IPTABLES_COMMENT}" ]] && commment="-m comment --comment \"$4\"" || comment=""
+: ${IPTABLES_CHAIN:=INPUT}
 
 function iptableAdd () {
     if [[ -z "${IPTABLES_INSERT}" ]]; then
-        iptables $comment -A INPUT -s $1 -j DROP
-        iptables $comment -A DOCKER-USER -s $1 -j DROP
+        iptables $comment -A $IPTABLES_CHAIN -s "$1" -j DROP
+
     else
-        iptables $comment -I INPUT $IPTABLES_INSERT -s $1 -j DROP
-        iptables $comment -I DOCKER-USER $IPTABLES_INSERT -s $1 -j DROP
+        iptables $comment -I $IPTABLES_CHAIN "$IPTABLES_INSERT" -s "$1" -j DROP
     fi
 }
 
 function iptableDel () {
-    iptables $comment -D INPUT -s $1 -j DROP
-    iptables $comment -D DOCKER-USER -s $1 -j DROP
+    iptables $comment -D $IPTABLES_CHAIN -s "$1" -j DROP
 }
 
 function ip6tableAdd () {
     if [[ -z "${IPTABLES_INSERT}" ]]; then
-        ip6tables $comment -A INPUT -s $1 -j DROP
-        ip6tables $comment -A DOCKER-USER -s $1 -j DROP
+        ip6tables $comment -A $IPTABLES_CHAIN -s "$1" -j DROP
     else
-        ip6tables $comment -I INPUT $IPTABLES_INSERT -s $1 -j DROP
-        ip6tables $comment -I DOCKER-USER $IPTABLES_INSERT -s $1 -j DROP
+        ip6tables $comment -I $IPTABLES_CHAIN "$IPTABLES_INSERT" -s "$1" -j DROP
     fi
 }
 
 function ip6tableDel () {
-    ip6tables $comment -D INPUT -s $1 -j DROP
-    ip6tables $comment -D DOCKER-USER -s $1 -j DROP
+    ip6tables $comment -D $IPTABLES_CHAIN -s "$1" -j DROP
 }
 
 #determine action
 if [ "$1" = "add" ]; then #add
     if [[ "$2" =~ .*[.].* ]]; then #ipv4
         echo "add $2 for $3 with $4"
-        iptableAdd $2
+        iptableAdd "$2"
     elif [[ "$2" =~ .*[:].* ]]; then #ipv6
         echo "IPV6 : add $2 for $3 with $4"
-        ip6tableAdd $2
+        ip6tableAdd "$2"
     fi
 elif [ "$1" = "del" ]; then #del
     if [[ "$2" =~ .*[.].* ]]; then #ipv4
         echo "del $2 for $3 with $4"
-        iptableDel $2
+        iptableDel "$2"
     elif [[ "$2" =~ .*[:].* ]]; then #ipv6
         echo "IPV6 : add $2 for $3 with $4"
-        ip6tableDel $2
+        ip6tableDel "$2"
     fi
 else
     echo  "unknon action"