From 275c7d54be2fe9acd963752a3857023dd44b0b9c Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Tue, 17 Dec 2024 19:54:59 +0000
Subject: [PATCH] Save rules backup after sync

---
 core/schains/firewall/firewall_manager.py |  5 +++++
 core/schains/firewall/iptables.py         |  3 +++
 core/schains/firewall/nftables.py         | 27 +++++++++++++++++++----
 core/schains/firewall/types.py            |  4 ++++
 tests/utils.py                            |  3 +++
 tools/configs/__init__.py                 |  2 ++
 6 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/core/schains/firewall/firewall_manager.py b/core/schains/firewall/firewall_manager.py
index d16f71574..1393864bc 100644
--- a/core/schains/firewall/firewall_manager.py
+++ b/core/schains/firewall/firewall_manager.py
@@ -71,6 +71,11 @@ def update_rules(self, rules: Iterable[SChainRule]) -> None:
         rules_to_remove = actual_rules - expected_rules
         self.add_rules(rules_to_add)
         self.remove_rules(rules_to_remove)
+        self.save_rules()
+
+    def save_rules(self) -> None:
+        """ Saves rules into persistent storage """
+        self.host_controller.save_rules()
 
     def add_rules(self, rules: Iterable[SChainRule]) -> None:
         logger.debug('Adding rules %s', rules)
diff --git a/core/schains/firewall/iptables.py b/core/schains/firewall/iptables.py
index 1d28c4037..fbe1b55f4 100644
--- a/core/schains/firewall/iptables.py
+++ b/core/schains/firewall/iptables.py
@@ -139,3 +139,6 @@ def from_ip_network(cls, ip: str) -> str:
     @classmethod
     def to_ip_network(cls, ip: str) -> str:
         return str(ipaddress.ip_network(ip))
+
+    def save_rules(self):
+        raise NotImplementedError('save_rules is not implemented for iptables host controller')
diff --git a/core/schains/firewall/nftables.py b/core/schains/firewall/nftables.py
index 669bd0273..57551ce5b 100644
--- a/core/schains/firewall/nftables.py
+++ b/core/schains/firewall/nftables.py
@@ -18,16 +18,17 @@
 #   along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 
-import logging
 import importlib
 import ipaddress
+import json
+import logging
 import multiprocessing
-from typing import Iterable
+import os
+from typing import Iterable, TypeVar
 
 from core.schains.firewall.types import IHostFirewallController, SChainRule
 
-from typing import TypeVar
-import json
+from tools.configs import NFT_CHAIN_BASE_PATH
 
 T = TypeVar('T')
 
@@ -315,3 +316,21 @@ def from_ip_network(cls, ip: str) -> str:
     @classmethod
     def to_ip_network(cls, ip: str) -> str:
         return str(ipaddress.ip_network(ip))
+
+    def get_plain_chain_rules(self) -> str:
+        self.nft.set_json_output(False)
+        output = ''
+        try:
+            rc, output, error = self.run_cmd(f'list chain {self.FAMILY} {self.table} {self.chain}')
+            if rc != 0:
+                raise NFTablesCmdFailedError(f"Failed to get table content: {error}")
+        finally:
+            self.nft.set_json_output(True)
+
+        return output
+
+    def save_rules(self) -> None:
+        chain_rules = self.get_plain_chain_rules()
+        nft_chain_path = os.path.join(NFT_CHAIN_BASE_PATH, f'{self.chain}.conf')
+        with open(nft_chain_path, 'w') as nft_chain_file:
+            nft_chain_file.write(chain_rules)
diff --git a/core/schains/firewall/types.py b/core/schains/firewall/types.py
index 65ba8885d..0d25e6d7f 100644
--- a/core/schains/firewall/types.py
+++ b/core/schains/firewall/types.py
@@ -88,6 +88,10 @@ def rules(self) -> Iterable[SChainRule]:  # pragma: no cover
     def has_rule(self, rule: SChainRule) -> bool:  # pragma: no cover
         pass
 
+    @abstractmethod
+    def save_rules(self) -> None:  # pragma: no cover
+        pass
+
 
 class IFirewallManager(ABC):
     @property
diff --git a/tests/utils.py b/tests/utils.py
index dc33bf91b..e7fa56881 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -222,6 +222,9 @@ def rules(self):
     def has_rule(self, srule):
         return srule in self._rules
 
+    def save_rules(self):
+        pass
+
 
 class SChainTestFirewallManager(SChainFirewallManager):
     def create_host_controller(self):
diff --git a/tools/configs/__init__.py b/tools/configs/__init__.py
index 4794de043..e1b0e053f 100644
--- a/tools/configs/__init__.py
+++ b/tools/configs/__init__.py
@@ -106,3 +106,5 @@
 SYNC_NODE = os.getenv('SYNC_NODE') == 'True'
 
 DOCKER_NODE_CONFIG_FILEPATH = os.path.join(NODE_DATA_PATH, 'docker.json')
+
+NFT_CHAIN_BASE_PATH = '/etc/nft.conf.d/chains'