From 3a8054d4d610d8e70fb982d5b38f21bef8b491d0 Mon Sep 17 00:00:00 2001
From: Tilman Kranz <tilman.kranz@b1-systems.de>
Date: Wed, 24 Nov 2021 01:57:30 +0100
Subject: [PATCH] add option "--config [CONFIG_FILE]"

---
 README.md        | 12 ++++----
 nft-edit-ruleset | 72 ++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 76 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 65ccb62..a01bc90 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,7 @@
 # Usage
 
 ```
-./nft-edit-ruleset [--yes] [--fail] [--timeout NUM]
+nft-edit-ruleset [--config [CONFIG_FILE]] [--fail] [--timeout SECONDS] [--yes]
 ```
 
 # Description
@@ -12,10 +12,12 @@ changes after a timeout given in seconds.
 
 # Options
 
-* `-f`, `--fail`: Exit unsuccessfully if changes fail to apply.
-* `-h`, `--help`: Display this message and exit.
-* `-t NUM`, `--timeout NUM`: Revert changes after NUM seconds.
-* `-y`, `--yes`: No confirmation before applying changes.
+- `-c`, `--config [CONFIG_FILE]`: On successfully applying the
+ changes, save ruleset to `CONFIG_FILE` (default: /etc/nftables.conf).
+- `-f`, `--fail`: Exit unsuccessfully if changes fail to apply.
+- `-h`, `--help`: Display this message and exit.
+- `-t NUM`, `--timeout NUM`: Revert changes after NUM seconds.
+- `-y`, `--yes`: No confirmation before applying changes.
 
 # Exit Codes
 
diff --git a/nft-edit-ruleset b/nft-edit-ruleset
index 32fa7f7..b0ef4bc 100755
--- a/nft-edit-ruleset
+++ b/nft-edit-ruleset
@@ -4,6 +4,17 @@
 # Authoremail: tilt@linuxfoo.de
 # License: MIT License (https://opensource.org/licenses/MIT)
 
+##
+# Configuration
+
+if test -w "/etc/nftables.conf" ; then
+  default_config_file="/etc/nftables.conf"
+elif test -w "/etc/sysconfig/nftables.conf" ; then
+  default_config_file="/etc/sysconfig/nftables.conf"
+else
+  default_config_file=""
+fi
+
 ##
 # Functions
 
@@ -33,23 +44,39 @@ set_timeout() {
     bash -c "nft -f '$ruleset' ; rm -f '$ruleset'"
 }
 
+store_config() {
+  local ruleset=$1
+  local conf=$2
+
+  if ! printf "#!/usr/sbin/nft -f\n# Generated at %s by $0\n\n" "$(date -R)" > "$conf" ; then
+    return 1
+  elif ! cat "$ruleset" >> "$conf" ; then
+    return 1
+  fi
+
+  return 0
+}
+
 ##
 # Arguments
 
-yes=false
+config=false
 fail=false
 timeout=false
+yes=false
 
 while true ; do
   case "$1" in
     -h|--help)
       cat << EOF
-Usage: $0 [--yes] [--fail] [--timeout NUM]
+Usage: nft-edit-ruleset [--config [CONFIG_FILE]] [--fail] [--timeout SECONDS] [--yes] 
 Description:
  Interactively edit the current nftables ruleset using the editor specified
  by environment variable EDITOR (defaulting to vim). Optionally, revert
  changes after a timeout given in seconds.
 Options:
+ - \`-c\`, \`--config [CONFIG_FILE]\`: On successfully applying the  
+  changes, save ruleset to \`CONFIG_FILE\` (default: $default_config_file).
  - \`-f\`, \`--fail\`: Exit unsuccessfully if changes fail to apply.
  - \`-h\`, \`--help\`: Display this message and exit.
  - \`-t NUM\`, \`--timeout NUM\`: Revert changes after NUM seconds.
@@ -70,9 +97,29 @@ EOF
         timeout_secs=$1
       else
         echo "ERROR: Missing argument for option \`--timeout\`, aborted."
-	exit 1
+      exit 1
       fi
       ;;
+    -c|--config)
+      config=true
+
+      if test "$#" -gt 0 && { test -w "$2" || test -w "$(dirname "$2")" ; } ; then
+        shift 1
+        config_file=$1
+      elif test -z "$default_config_file" ; then
+        echo "ERROR: Option \`--config\` was used without an argument, but no default location of a file \"nftables.conf\" could be found (use \`--config CONFIG_FILE\` to specify a location); aborted."
+        exit 1
+      else
+        config_file=$default_config_file
+      fi
+      ;;
+    '')
+      :
+      ;;
+    *)
+      echo "ERROR: Unknown or unexpected argument \"$1\"; aborted."
+      exit 1
+      ;;
   esac
 
   if test "$#" -gt 0 ; then
@@ -109,6 +156,16 @@ while true ; do
 
   if diff "$tmp_old" "$tmp" ; then
     echo "No changes."
+
+    if "$config" ; then
+      if store_config "$tmp" "$config_file" ; then
+        echo "Stored unchanged ruleset to config_file=\"$config_file\"."
+      else
+        echo "ERROR: Storing unchanged ruleset to config_file=\"$config_file\" failed."
+        exit 1
+      fi
+    fi
+
     exit 0
   else
     if ! "$yes" ; then
@@ -144,6 +201,15 @@ while true ; do
 
       echo "Changes applied successfully."
 
+      if "$config" ; then
+        if store_config "$tmp" "$config_file" ; then
+          echo "Stored changed ruleset to config_file=\"$config_file\"."
+        else
+          echo "ERROR: Storing changed ruleset to config_file=\"$config_file\" failed."
+          exit 1
+        fi
+      fi
+
       break
     fi
   fi