From c657e0c64ef7e73eb3de833b012322dc528426bc Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Tue, 14 Jan 2025 16:47:17 +0000
Subject: [PATCH 01/11] Fix nftables chain cleanup

---
 core/schains/firewall/nftables.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/core/schains/firewall/nftables.py b/core/schains/firewall/nftables.py
index 00d81516..486846cf 100644
--- a/core/schains/firewall/nftables.py
+++ b/core/schains/firewall/nftables.py
@@ -23,6 +23,7 @@
 import logging
 import multiprocessing
 import os
+import shutil
 from typing import Iterable
 
 from core.schains.firewall.types import IHostFirewallController, SChainRule
@@ -343,5 +344,10 @@ def save_rules(self) -> None:
         with open(nft_chain_path, 'w') as nft_chain_file:
             nft_chain_file.write(chain_rules)
 
+    def remove_saved_rules(self) -> None:
+        nft_chain_path = os.path.join(NFT_CHAIN_BASE_PATH, f'{self.chain}.conf')
+        shutil.rmtree(nft_chain_path)
+
     def cleanup(self) -> None:
         self.delete_chain()
+        self.remove_saved_rules()

From 934cee45748355a327c34dc405e49cff7607c97f Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Wed, 15 Jan 2025 13:37:51 +0000
Subject: [PATCH 02/11] Make cleaner to remove schain firewall rules config

---
 core/schains/firewall/nftables.py |  8 +++++--
 tests/firewall/nftables_test.py   | 36 ++++++++++++++++++++++++++++++-
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/core/schains/firewall/nftables.py b/core/schains/firewall/nftables.py
index 486846cf..b7b31260 100644
--- a/core/schains/firewall/nftables.py
+++ b/core/schains/firewall/nftables.py
@@ -52,6 +52,7 @@ def __init__(self, table: str = TABLE, chain: str = CHAIN) -> None:
         self._nftables = importlib.import_module('nftables')
         self.nft = self._nftables.Nftables()
         self.nft.set_json_output(True)
+        self.nft.set_stateless_output(True)
 
     @classmethod
     def rule_to_expr(cls, rule: SChainRule, counter: bool = True) -> list:
@@ -179,6 +180,7 @@ def create_chain(self, first_port: int, last_port: int) -> None:
                 )
             )
             self.add_schain_drop_rule(first_port, last_port)
+            self.save_rules()
 
     def delete_chain(self) -> None:
         if self.has_chain(self.chain):
@@ -330,7 +332,9 @@ 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}')
+            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:
@@ -346,7 +350,7 @@ def save_rules(self) -> None:
 
     def remove_saved_rules(self) -> None:
         nft_chain_path = os.path.join(NFT_CHAIN_BASE_PATH, f'{self.chain}.conf')
-        shutil.rmtree(nft_chain_path)
+        os.remove(nft_chain_path)
 
     def cleanup(self) -> None:
         self.delete_chain()
diff --git a/tests/firewall/nftables_test.py b/tests/firewall/nftables_test.py
index 4481265f..c030aa26 100644
--- a/tests/firewall/nftables_test.py
+++ b/tests/firewall/nftables_test.py
@@ -1,11 +1,14 @@
 import concurrent.futures
 import importlib
+import os
+import shutil
 import time
 
 import pytest
 
-from core.schains.firewall.nftables import NFTablesController
+from core.schains.firewall.nftables import NFTablesController, NFT_CHAIN_BASE_PATH
 from core.schains.firewall.types import SChainRule
+from tools.helper import run_cmd
 
 
 @pytest.fixture
@@ -15,6 +18,16 @@ def nf_test_tables():
     return nft
 
 
+@pytest.fixture()
+def nft_chain_folder():
+    path = '/etc/nft.conf.d/chains'
+    try:
+        os.makedirs(path)
+        yield path
+    finally:
+        shutil.rmtree(path)
+
+
 @pytest.fixture
 def filter_table(nf_test_tables):
     print(nf_test_tables.cmd('add table inet firewall'))
@@ -67,6 +80,27 @@ def test_nftables_controller_duplicates(custom_chain):
     ]
 
 
+def test_create_delete_chain(filter_table, nft_chain_folder):
+    chain_name = 'test-chain'
+
+    output = run_cmd(['nft', 'list', 'chains']).stdout.decode('utf-8')
+    output == 'table inet firewall {\n}\n'
+    nft_chain_path = os.path.join(NFT_CHAIN_BASE_PATH, f'skale-{chain_name}.conf')
+    assert not os.path.isfile(nft_chain_path)
+
+    manager = NFTablesController(chain=chain_name)
+    manager.create_chain(first_port=10000, last_port=10063)
+
+    chains = run_cmd(['nft', 'list', 'chains']).stdout.decode('utf-8')
+    assert chains == 'table inet firewall {\n\tchain skale-test-chain {\n\t\ttype filter hook input priority filter; policy accept;\n\t}\n}\n'  # noqa
+    assert os.path.isfile(nft_chain_path)
+
+    manager.cleanup()
+    chains = run_cmd(['nft', 'list', 'chains']).stdout.decode('utf-8')
+    assert chains == 'table inet firewall {\n}\n'
+    assert not os.path.isfile(nft_chain_path)
+
+
 def add_remove_rule(srule, refresh):
     manager = NFTablesController()
     manager.add_rule(srule)

From 4a1c830892f29023642bd866da4dc7d359e5621f Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Wed, 15 Jan 2025 13:38:53 +0000
Subject: [PATCH 03/11] Fix linter

---
 core/schains/firewall/nftables.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/core/schains/firewall/nftables.py b/core/schains/firewall/nftables.py
index b7b31260..f333dd4f 100644
--- a/core/schains/firewall/nftables.py
+++ b/core/schains/firewall/nftables.py
@@ -23,7 +23,6 @@
 import logging
 import multiprocessing
 import os
-import shutil
 from typing import Iterable
 
 from core.schains.firewall.types import IHostFirewallController, SChainRule

From 218ede8ec932d44217fc527444d04a2be4ea1d16 Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Wed, 15 Jan 2025 17:53:57 +0000
Subject: [PATCH 04/11] Reorder operations

---
 core/schains/firewall/nftables.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/schains/firewall/nftables.py b/core/schains/firewall/nftables.py
index f333dd4f..d94d9162 100644
--- a/core/schains/firewall/nftables.py
+++ b/core/schains/firewall/nftables.py
@@ -352,5 +352,5 @@ def remove_saved_rules(self) -> None:
         os.remove(nft_chain_path)
 
     def cleanup(self) -> None:
-        self.delete_chain()
         self.remove_saved_rules()
+        self.delete_chain()

From f4361cb9f4294e8756f8efb208e7946cbdfc7623 Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Sat, 18 Jan 2025 16:03:32 +0000
Subject: [PATCH 05/11] Make sure rules persistent and chain configured

---
 core/schains/checks.py                        | 18 +++++++-
 core/schains/firewall/firewall_manager.py     |  8 ++++
 core/schains/firewall/nftables.py             | 44 +++++++++++++++----
 core/schains/firewall/rule_controller.py      | 16 +++++++
 core/schains/firewall/types.py                |  8 ++++
 tests/conftest.py                             | 10 +++++
 .../firewall/default_rule_controller_test.py  | 12 -----
 tests/firewall/nftables_test.py               | 11 -----
 tests/utils.py                                |  6 +++
 tools/configs/__init__.py                     |  2 +-
 10 files changed, 100 insertions(+), 35 deletions(-)

diff --git a/core/schains/checks.py b/core/schains/checks.py
index 8beb2673..c56a8d72 100644
--- a/core/schains/checks.py
+++ b/core/schains/checks.py
@@ -301,6 +301,12 @@ def volume(self) -> CheckRes:
     @property
     def firewall_rules(self) -> CheckRes:
         """Checks that firewall rules are set correctly"""
+        data = {
+            'config': False,
+            'inited': False,
+            'rules': False,
+            'persistant': False,
+        }
         if self.config:
             conf = self.cfm.skaled_config
             base_port = get_base_port_from_config(conf)
@@ -311,8 +317,16 @@ def firewall_rules(self) -> CheckRes:
                 base_port=base_port, own_ip=own_ip, node_ips=node_ips, sync_ip_ranges=ranges
             )
             logger.debug(f'Rule controller {self.rc.expected_rules()}')
-            return CheckRes(self.rc.is_rules_synced())
-        return CheckRes(False)
+            data = {
+                'config': True,
+                'inited': self.rc.is_inited(),
+                'rules': self.rc.is_rules_synced(),
+                'persistent': self.rc.is_persistent(),
+            }
+            logger.debug('Firewall rules check: %s', data)
+            status = all(data.values())
+            return CheckRes(status=status, data=data)
+        return CheckRes(status=False, data=data)
 
     @property
     def skaled_container(self) -> CheckRes:
diff --git a/core/schains/firewall/firewall_manager.py b/core/schains/firewall/firewall_manager.py
index 5b2f7640..6c9f3768 100644
--- a/core/schains/firewall/firewall_manager.py
+++ b/core/schains/firewall/firewall_manager.py
@@ -103,3 +103,11 @@ def create_host_controller(self) -> NFTablesController:
         nc_controller.create_table()
         nc_controller.create_chain(self.first_port, self.last_port)
         return nc_controller
+
+    def rules_saved(self) -> bool:
+        saved = self.host_controller.get_saved_rules()
+        return saved == self.host_controller.get_rules()
+
+    def base_config_applied(self) -> bool:
+        return self.host_controller.has_chain(self.host_controller.chain) and \
+            self.host_controller.has_drop_rule(self.first_port, self.last_port)
diff --git a/core/schains/firewall/nftables.py b/core/schains/firewall/nftables.py
index d94d9162..04f2c753 100644
--- a/core/schains/firewall/nftables.py
+++ b/core/schains/firewall/nftables.py
@@ -1,5 +1,4 @@
 #   -*- coding: utf-8 -*-
-#
 #   This file is part of SKALE Admin
 #
 #   Copyright (C) 2024 SKALE Labs
@@ -125,7 +124,7 @@ def create_table(self) -> None:
         if not self.has_table(self.table):
             return self.run_cmd(f'add table inet {self.table}')
 
-    def add_schain_drop_rule(self, first_port: int, last_port: int) -> None:
+    def has_drop_rule(self, first_port: int, last_port: int) -> bool:
         expr = [
             {
                 'match': {
@@ -138,7 +137,22 @@ def add_schain_drop_rule(self, first_port: int, last_port: int) -> None:
             {'drop': None},
         ]
 
-        if self.expr_to_rule(expr) not in self.get_rules_by_policy(policy='drop'):
+        return self.expr_to_rule(expr) in self.get_rules_by_policy(policy='drop')
+
+    def add_schain_drop_rule(self, first_port: int, last_port: int) -> None:
+        if not self.has_drop_rule(first_port, last_port):
+            expr = [
+                {
+                    'match': {
+                        'op': '==',
+                        'left': {'payload': {'protocol': 'tcp', 'field': 'dport'}},
+                        'right': {'range': [first_port, last_port]},
+                    }
+                },
+                {'counter': None},
+                {'drop': None},
+            ]
+
             cmd = {
                 'nftables': [
                     {
@@ -178,8 +192,8 @@ def create_chain(self, first_port: int, last_port: int) -> None:
                     ]
                 )
             )
-            self.add_schain_drop_rule(first_port, last_port)
-            self.save_rules()
+        self.add_schain_drop_rule(first_port, last_port)
+        self.save_rules()
 
     def delete_chain(self) -> None:
         if self.has_chain(self.chain):
@@ -339,17 +353,29 @@ def get_plain_chain_rules(self) -> str:
         finally:
             self.nft.set_json_output(True)
 
+        # cleanup table header
+        output = '\n'.join(output.split('\n')[2:-1])
+
         return output
 
+    @property
+    def nft_chain_path(self) -> str:
+        return os.path.join(NFT_CHAIN_BASE_PATH, f'{self.chain}.conf')
+
     def save_rules(self) -> None:
+        logger.info('Saving the firewall rules for chain %s', self.chain)
         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:
+        with open(self.nft_chain_path, 'w') as nft_chain_file:
             nft_chain_file.write(chain_rules)
 
+    def get_saved_rules(self) -> str:
+        if not os.path.isfile(self.nft_chain_path):
+            return ''
+        with open(self.nft_chain_path, 'r') as nft_chain_file:
+            return nft_chain_file.read()
+
     def remove_saved_rules(self) -> None:
-        nft_chain_path = os.path.join(NFT_CHAIN_BASE_PATH, f'{self.chain}.conf')
-        os.remove(nft_chain_path)
+        os.remove(self.nft_chain_path)
 
     def cleanup(self) -> None:
         self.remove_saved_rules()
diff --git a/core/schains/firewall/rule_controller.py b/core/schains/firewall/rule_controller.py
index 3b63026b..a109d30b 100644
--- a/core/schains/firewall/rule_controller.py
+++ b/core/schains/firewall/rule_controller.py
@@ -215,6 +215,14 @@ def create_firewall_manager(self) -> IptablesSChainFirewallManager:
             self.base_port + self.ports_per_schain - 1  # type: ignore
         )
 
+    @configured_only
+    def is_persistent(self) -> bool:
+        return True
+
+    @configured_only
+    def is_inited(self) -> bool:
+        return True
+
 
 class NFTSchainRuleController(SChainRuleController):
     @configured_only
@@ -224,3 +232,11 @@ def create_firewall_manager(self) -> NFTSchainFirewallManager:
             self.base_port,  # type: ignore
             self.base_port + self.ports_per_schain - 1  # type: ignore
         )
+
+    @configured_only
+    def is_persistent(self) -> bool:
+        return self.firewall_manager.rules_saved()
+
+    @configured_only
+    def is_inited(self) -> bool:
+        return self.firewall_manager.base_config_applied()
diff --git a/core/schains/firewall/types.py b/core/schains/firewall/types.py
index 0062ccce..ecb076c6 100644
--- a/core/schains/firewall/types.py
+++ b/core/schains/firewall/types.py
@@ -139,3 +139,11 @@ def sync(self) -> None:  # pragma: no cover
     @abstractmethod
     def cleanup(self) -> None:  # pragma: no cover
         pass
+
+    @abstractmethod
+    def is_persistent(self) -> bool:  # pragma: no cover
+        pass
+
+    @abstractmethod
+    def is_inited(self) -> bool:  # pragma: no cover
+        pass
diff --git a/tests/conftest.py b/tests/conftest.py
index 973a375e..ac4b92ea 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -631,3 +631,13 @@ def ncli_status(_schain_name):
         yield init_node_cli_status(_schain_name)
     finally:
         shutil.rmtree(schain_dir_path, ignore_errors=True)
+
+
+@pytest.fixture()
+def nft_chain_folder():
+    path = '/etc/nft.conf.d/skale/chains'
+    try:
+        os.makedirs(path)
+        yield path
+    finally:
+        shutil.rmtree(path)
diff --git a/tests/firewall/default_rule_controller_test.py b/tests/firewall/default_rule_controller_test.py
index c21489a8..cadd5f5d 100644
--- a/tests/firewall/default_rule_controller_test.py
+++ b/tests/firewall/default_rule_controller_test.py
@@ -1,8 +1,6 @@
 
 import concurrent.futures
 import mock
-import os
-import shutil
 
 import pytest
 
@@ -24,16 +22,6 @@ def refresh():
         run_cmd(['nft', 'flush', 'ruleset'])
 
 
-@pytest.fixture()
-def nft_chain_folder():
-    path = '/etc/nft.conf.d/chains'
-    try:
-        os.makedirs(path)
-        yield path
-    finally:
-        shutil.rmtree(path)
-
-
 def test_get_default_rule_controller(nft_chain_folder):
     own_ip = '3.3.3.3'
     node_ips = ['1.1.1.1', '2.2.2.2', '3.3.3.3', '4.4.4.4']
diff --git a/tests/firewall/nftables_test.py b/tests/firewall/nftables_test.py
index c030aa26..9f7f8c63 100644
--- a/tests/firewall/nftables_test.py
+++ b/tests/firewall/nftables_test.py
@@ -1,7 +1,6 @@
 import concurrent.futures
 import importlib
 import os
-import shutil
 import time
 
 import pytest
@@ -18,16 +17,6 @@ def nf_test_tables():
     return nft
 
 
-@pytest.fixture()
-def nft_chain_folder():
-    path = '/etc/nft.conf.d/chains'
-    try:
-        os.makedirs(path)
-        yield path
-    finally:
-        shutil.rmtree(path)
-
-
 @pytest.fixture
 def filter_table(nf_test_tables):
     print(nf_test_tables.cmd('add table inet firewall'))
diff --git a/tests/utils.py b/tests/utils.py
index 014c054f..06b7e515 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -242,6 +242,12 @@ def create_firewall_manager(self):
             self.base_port + self.ports_per_schain
         )
 
+    def is_persistent(self) -> bool:
+        return True
+
+    def is_inited(self) -> bool:
+        return True
+
 
 def get_test_rule_controller(
     name,
diff --git a/tools/configs/__init__.py b/tools/configs/__init__.py
index e1b0e053..ec04f617 100644
--- a/tools/configs/__init__.py
+++ b/tools/configs/__init__.py
@@ -107,4 +107,4 @@
 
 DOCKER_NODE_CONFIG_FILEPATH = os.path.join(NODE_DATA_PATH, 'docker.json')
 
-NFT_CHAIN_BASE_PATH = '/etc/nft.conf.d/chains'
+NFT_CHAIN_BASE_PATH = '/etc/nft.conf.d/skale/chains'

From 8bacef041dc7eb0438c2750b41b52fe27cc3027d Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Tue, 21 Jan 2025 12:09:12 +0000
Subject: [PATCH 06/11] Save firewall rules properly

---
 core/schains/firewall/firewall_manager.py |  2 +-
 core/schains/firewall/nftables.py         | 13 +++++++++++--
 tests/firewall/nftables_test.py           | 16 ++++++++++++++++
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/core/schains/firewall/firewall_manager.py b/core/schains/firewall/firewall_manager.py
index 6c9f3768..6b33ab3b 100644
--- a/core/schains/firewall/firewall_manager.py
+++ b/core/schains/firewall/firewall_manager.py
@@ -106,7 +106,7 @@ def create_host_controller(self) -> NFTablesController:
 
     def rules_saved(self) -> bool:
         saved = self.host_controller.get_saved_rules()
-        return saved == self.host_controller.get_rules()
+        return saved == self.host_controller.get_plain_chain_rules()
 
     def base_config_applied(self) -> bool:
         return self.host_controller.has_chain(self.host_controller.chain) and \
diff --git a/core/schains/firewall/nftables.py b/core/schains/firewall/nftables.py
index 04f2c753..82addc75 100644
--- a/core/schains/firewall/nftables.py
+++ b/core/schains/firewall/nftables.py
@@ -353,9 +353,18 @@ def get_plain_chain_rules(self) -> str:
         finally:
             self.nft.set_json_output(True)
 
+        lines = output.split('\n')
         # cleanup table header
-        output = '\n'.join(output.split('\n')[2:-1])
-
+        if lines[-1] == '':
+            lines = lines[1:-2]
+        else:
+            lines = lines[1:-1]
+
+        # remove leading tab
+        lines = list(map(lambda line: line[1:], lines))
+        # Adding new line at the end to prevent validation failure
+        lines.append('')
+        output = '\n'.join(lines)
         return output
 
     @property
diff --git a/tests/firewall/nftables_test.py b/tests/firewall/nftables_test.py
index 9f7f8c63..cee70eed 100644
--- a/tests/firewall/nftables_test.py
+++ b/tests/firewall/nftables_test.py
@@ -90,6 +90,22 @@ def test_create_delete_chain(filter_table, nft_chain_folder):
     assert not os.path.isfile(nft_chain_path)
 
 
+def test_saved_rules(filter_table, nft_chain_folder):
+    chain_name = 'test-chain'
+    nft_chain_path = os.path.join(NFT_CHAIN_BASE_PATH, f'skale-{chain_name}.conf')
+
+    manager = NFTablesController(chain=chain_name)
+    assert not os.path.isfile(nft_chain_path)
+    manager.create_chain(first_port=10000, last_port=10063)
+    assert os.path.isfile(nft_chain_path)
+    assert manager.get_saved_rules() == 'chain skale-test-chain {\n\ttype filter hook input priority filter; policy accept;\n\ttcp dport 10000-10063 counter drop\n}\n'  # noqa
+
+    assert os.path.isfile(nft_chain_path)
+
+    manager.remove_saved_rules()
+    assert not os.path.isfile(nft_chain_path)
+
+
 def add_remove_rule(srule, refresh):
     manager = NFTablesController()
     manager.add_rule(srule)

From 6ce86c086d3e3fa034696e752e5ec07fc61a6022 Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Tue, 21 Jan 2025 20:12:46 +0000
Subject: [PATCH 07/11] Fix cleaner procedure

---
 core/schains/checks.py            |  2 --
 core/schains/cleaner.py           | 30 ++++++++++++++++++++++++++----
 core/schains/firewall/nftables.py |  3 ++-
 tools/configs/__init__.py         |  1 +
 4 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/core/schains/checks.py b/core/schains/checks.py
index c56a8d72..0a42bbdb 100644
--- a/core/schains/checks.py
+++ b/core/schains/checks.py
@@ -302,7 +302,6 @@ def volume(self) -> CheckRes:
     def firewall_rules(self) -> CheckRes:
         """Checks that firewall rules are set correctly"""
         data = {
-            'config': False,
             'inited': False,
             'rules': False,
             'persistant': False,
@@ -318,7 +317,6 @@ def firewall_rules(self) -> CheckRes:
             )
             logger.debug(f'Rule controller {self.rc.expected_rules()}')
             data = {
-                'config': True,
                 'inited': self.rc.is_inited(),
                 'rules': self.rc.is_rules_synced(),
                 'persistent': self.rc.is_persistent(),
diff --git a/core/schains/cleaner.py b/core/schains/cleaner.py
index 7fd291ef..0511e4e7 100644
--- a/core/schains/cleaner.py
+++ b/core/schains/cleaner.py
@@ -17,6 +17,7 @@
 #   You should have received a copy of the GNU Affero General Public License
 #   along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
+import glob
 import logging
 import os
 import shutil
@@ -43,7 +44,11 @@
 from core.schains.types import ContainerType
 from core.schains.firewall.utils import get_sync_agent_ranges
 
-from tools.configs import SGX_CERTIFICATES_FOLDER, SYNC_NODE
+from tools.configs import (
+    NFT_CHAIN_CONFIG_WILDCARD,
+    SGX_CERTIFICATES_FOLDER,
+    SYNC_NODE
+)
 from tools.configs.schains import SCHAINS_DIR_PATH
 from tools.configs.containers import SCHAIN_CONTAINER, IMA_CONTAINER, SCHAIN_STOP_TIMEOUT
 from tools.docker_utils import DockerUtils
@@ -136,18 +141,34 @@ def get_schains_with_containers(dutils=None):
     ]
 
 
+def get_schains_firewall_configs() -> list:
+    return list(map(lambda path: os.path.basename(path), glob.glob(NFT_CHAIN_CONFIG_WILDCARD)))
+
+
 def get_schains_on_node(dutils=None):
     dutils = dutils or DockerUtils()
     schains_with_dirs = os.listdir(SCHAINS_DIR_PATH)
     schains_with_container = get_schains_with_containers(dutils)
     schains_active_records = get_schains_names()
+    schains_firewall_configs = list(
+        map(lambda name: name.removeprefix('skale-'),
+        get_schains_firewall_configs())
+    )
     logger.info(
-        'dirs %s, containers: %s, records: %s',
+        'dirs %s, containers: %s, records: %s, firewall configs: %s',
         schains_with_dirs,
         schains_with_container,
-        schains_active_records
+        schains_active_records,
+        schains_firewall_configs
+    )
+    return sorted(
+        merged_unique(
+            schains_with_dirs,
+            schains_with_container,
+            schains_active_records,
+            schains_firewall_configs
+        )
     )
-    return sorted(merged_unique(schains_with_dirs, schains_with_container, schains_active_records))
 
 
 def schain_names_to_ids(skale, schain_names):
@@ -268,6 +289,7 @@ def cleanup_schain(
             ranges = estate.ranges
         rc.configure(base_port=base_port, own_ip=own_ip, node_ips=node_ips, sync_ip_ranges=ranges)
         rc.cleanup()
+
     if estate is not None and estate.ima_linked:
         if check_status.get('ima_container', False) or is_exited(
             schain_name, container_type=ContainerType.ima, dutils=dutils
diff --git a/core/schains/firewall/nftables.py b/core/schains/firewall/nftables.py
index 82addc75..a2977c4b 100644
--- a/core/schains/firewall/nftables.py
+++ b/core/schains/firewall/nftables.py
@@ -384,7 +384,8 @@ def get_saved_rules(self) -> str:
             return nft_chain_file.read()
 
     def remove_saved_rules(self) -> None:
-        os.remove(self.nft_chain_path)
+        if os.isfile(self.nft_chain_path):
+            os.remove(self.nft_chain_path)
 
     def cleanup(self) -> None:
         self.remove_saved_rules()
diff --git a/tools/configs/__init__.py b/tools/configs/__init__.py
index ec04f617..1e2fd423 100644
--- a/tools/configs/__init__.py
+++ b/tools/configs/__init__.py
@@ -108,3 +108,4 @@
 DOCKER_NODE_CONFIG_FILEPATH = os.path.join(NODE_DATA_PATH, 'docker.json')
 
 NFT_CHAIN_BASE_PATH = '/etc/nft.conf.d/skale/chains'
+NFT_CHAIN_CONFIG_WILDCARD = os.path.join(NFT_CHAIN_BASE_PATH, '*')

From 241ab43f0b392d146ff11c1209fa01e5423679bd Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Tue, 21 Jan 2025 21:30:40 +0000
Subject: [PATCH 08/11] Fix glob

---
 core/schains/cleaner.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/core/schains/cleaner.py b/core/schains/cleaner.py
index 0511e4e7..9d1bd1f4 100644
--- a/core/schains/cleaner.py
+++ b/core/schains/cleaner.py
@@ -22,6 +22,7 @@
 import os
 import shutil
 from multiprocessing import Process
+from pathlib import Path
 from typing import Optional
 
 from sgx import SgxClient
@@ -142,7 +143,7 @@ def get_schains_with_containers(dutils=None):
 
 
 def get_schains_firewall_configs() -> list:
-    return list(map(lambda path: os.path.basename(path), glob.glob(NFT_CHAIN_CONFIG_WILDCARD)))
+    return list(map(lambda path: Path(path).stem, glob.glob(NFT_CHAIN_CONFIG_WILDCARD)))
 
 
 def get_schains_on_node(dutils=None):
@@ -279,7 +280,7 @@ def cleanup_schain(
         remove_schain_container(schain_name, dutils=dutils)
     if check_status['volume']:
         remove_schain_volume(schain_name, dutils=dutils)
-    if check_status['firewall_rules']:
+    if any(checks.firewall_rules.data):
         conf = ConfigFileManager(schain_name).skaled_config
         base_port = get_base_port_from_config(conf)
         own_ip = get_own_ip_from_config(conf)

From 9e33f8278e0df0ae598563771d0a327e538a70f4 Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Wed, 22 Jan 2025 13:58:55 +0000
Subject: [PATCH 09/11] Handle firewall related cleaner failure

---
 core/schains/cleaner.py                   | 25 +++++++----------------
 core/schains/firewall/firewall_manager.py | 11 ++++++----
 core/schains/firewall/nftables.py         |  5 ++---
 core/schains/firewall/rule_controller.py  | 10 ++++++---
 core/schains/firewall/types.py            |  2 +-
 core/schains/firewall/utils.py            |  7 +++++++
 tests/firewall/firewall_manager_test.py   |  4 ++--
 tests/firewall/nftables_test.py           | 19 ++++++++++++++++-
 tests/schains/cleaner_test.py             |  5 ++++-
 tests/utils.py                            |  6 ++++++
 10 files changed, 61 insertions(+), 33 deletions(-)

diff --git a/core/schains/cleaner.py b/core/schains/cleaner.py
index 9d1bd1f4..881a2e63 100644
--- a/core/schains/cleaner.py
+++ b/core/schains/cleaner.py
@@ -30,15 +30,9 @@
 
 from core.node import get_current_nodes, get_skale_node_version
 from core.schains.checks import SChainChecks
-from core.schains.config.file_manager import ConfigFileManager
 from core.schains.config.directory import schain_config_dir
 from core.schains.dkg.utils import get_secret_key_share_filepath
-from core.schains.firewall.utils import get_default_rule_controller
-from core.schains.config.helper import (
-    get_base_port_from_config,
-    get_node_ips_from_config,
-    get_own_ip_from_config,
-)
+from core.schains.firewall.utils import cleanup_firewall_for_schain, get_default_rule_controller
 from core.schains.process import ProcessReport, terminate_process
 from core.schains.runner import get_container_name, is_exited
 from core.schains.external_config import ExternalConfig
@@ -152,8 +146,10 @@ def get_schains_on_node(dutils=None):
     schains_with_container = get_schains_with_containers(dutils)
     schains_active_records = get_schains_names()
     schains_firewall_configs = list(
-        map(lambda name: name.removeprefix('skale-'),
-        get_schains_firewall_configs())
+        map(
+            lambda name: name.removeprefix('skale-'),
+            get_schains_firewall_configs()
+        )
     )
     logger.info(
         'dirs %s, containers: %s, records: %s, firewall configs: %s',
@@ -281,15 +277,8 @@ def cleanup_schain(
     if check_status['volume']:
         remove_schain_volume(schain_name, dutils=dutils)
     if any(checks.firewall_rules.data):
-        conf = ConfigFileManager(schain_name).skaled_config
-        base_port = get_base_port_from_config(conf)
-        own_ip = get_own_ip_from_config(conf)
-        node_ips = get_node_ips_from_config(conf)
-        ranges = []
-        if estate is not None:
-            ranges = estate.ranges
-        rc.configure(base_port=base_port, own_ip=own_ip, node_ips=node_ips, sync_ip_ranges=ranges)
-        rc.cleanup()
+        logger.info('Cleaning firewall for %s', schain_name)
+        cleanup_firewall_for_schain(schain_name)
 
     if estate is not None and estate.ima_linked:
         if check_status.get('ima_container', False) or is_exited(
diff --git a/core/schains/firewall/firewall_manager.py b/core/schains/firewall/firewall_manager.py
index 6b33ab3b..f9d1bad2 100644
--- a/core/schains/firewall/firewall_manager.py
+++ b/core/schains/firewall/firewall_manager.py
@@ -87,15 +87,14 @@ def remove_rules(self, rules: Iterable[SChainRule]) -> None:
         for rule in rules:
             self.host_controller.remove_rule(rule)
 
-    def flush(self) -> None:
-        self.remove_rules(self.rules)
-        self.host_controller.cleanup()
-
 
 class IptablesSChainFirewallManager(SChainFirewallManager):
     def create_host_controller(self) -> IptablesController:
         return IptablesController()
 
+    def cleanup(self) -> None:
+        self.remove_rules(self.rules)
+
 
 class NFTSchainFirewallManager(SChainFirewallManager):
     def create_host_controller(self) -> NFTablesController:
@@ -111,3 +110,7 @@ def rules_saved(self) -> bool:
     def base_config_applied(self) -> bool:
         return self.host_controller.has_chain(self.host_controller.chain) and \
             self.host_controller.has_drop_rule(self.first_port, self.last_port)
+
+    def cleanup(self) -> None:
+        self.host_controller.cleanup()
+        self.host_controller.remove_saved_rules()
diff --git a/core/schains/firewall/nftables.py b/core/schains/firewall/nftables.py
index a2977c4b..e2a56e39 100644
--- a/core/schains/firewall/nftables.py
+++ b/core/schains/firewall/nftables.py
@@ -44,7 +44,7 @@ class NFTablesController(IHostFirewallController):
     plock = multiprocessing.Lock()
     FAMILY = 'inet'
 
-    def __init__(self, table: str = TABLE, chain: str = CHAIN) -> None:
+    def __init__(self, chain: str, table: str = TABLE) -> None:
         self.table = table
         self.chain = f'skale-{chain}'
         self._nftables = importlib.import_module('nftables')
@@ -384,9 +384,8 @@ def get_saved_rules(self) -> str:
             return nft_chain_file.read()
 
     def remove_saved_rules(self) -> None:
-        if os.isfile(self.nft_chain_path):
+        if os.path.isfile(self.nft_chain_path):
             os.remove(self.nft_chain_path)
 
     def cleanup(self) -> None:
-        self.remove_saved_rules()
         self.delete_chain()
diff --git a/core/schains/firewall/rule_controller.py b/core/schains/firewall/rule_controller.py
index a109d30b..68620526 100644
--- a/core/schains/firewall/rule_controller.py
+++ b/core/schains/firewall/rule_controller.py
@@ -202,9 +202,6 @@ def sync(self) -> None:
         logger.debug('Syncing firewall rules with %s', erules)
         self.firewall_manager.update_rules(erules)
 
-    def cleanup(self) -> None:
-        self.firewall_manager.flush()
-
 
 class IptablesSChainRuleController(SChainRuleController):
     @configured_only
@@ -223,6 +220,10 @@ def is_persistent(self) -> bool:
     def is_inited(self) -> bool:
         return True
 
+    @configured_only
+    def cleanup(self) -> None:
+        self.firewall_manager.cleanup()
+
 
 class NFTSchainRuleController(SChainRuleController):
     @configured_only
@@ -240,3 +241,6 @@ def is_persistent(self) -> bool:
     @configured_only
     def is_inited(self) -> bool:
         return self.firewall_manager.base_config_applied()
+
+    def cleanup(self) -> None:
+        self.firewall_manager.cleanup()
diff --git a/core/schains/firewall/types.py b/core/schains/firewall/types.py
index ecb076c6..c30bfc11 100644
--- a/core/schains/firewall/types.py
+++ b/core/schains/firewall/types.py
@@ -108,7 +108,7 @@ def update_rules(self, rules: Iterable[SChainRule]) -> None:  # pragma: no cover
         pass
 
     @abstractmethod
-    def flush(self) -> None:  # pragma: no cover  # noqa
+    def cleanup(self) -> None:  # pragma: no cover  # noqa
         pass
 
 
diff --git a/core/schains/firewall/utils.py b/core/schains/firewall/utils.py
index 1f94694f..0788c6df 100644
--- a/core/schains/firewall/utils.py
+++ b/core/schains/firewall/utils.py
@@ -25,6 +25,7 @@
 from skale import Skale
 
 from .types import IpRange
+from .nftables import NFTablesController
 from .rule_controller import IptablesSChainRuleController, NFTSchainRuleController
 
 
@@ -101,3 +102,9 @@ def save_sync_ranges(sync_agent_ranges: List[IpRange], path: str) -> None:
 
 def ranges_from_plain_tuples(plain_ranges: List[Tuple]) -> List[IpRange]:
     return list(sorted(map(lambda r: IpRange(*r), plain_ranges)))
+
+
+def cleanup_firewall_for_schain(schain_name: str) -> None:
+    nft = NFTablesController(chain=schain_name)
+    nft.cleanup()
+    nft.remove_saved_rules()
diff --git a/tests/firewall/firewall_manager_test.py b/tests/firewall/firewall_manager_test.py
index 719ad1bf..04203acc 100644
--- a/tests/firewall/firewall_manager_test.py
+++ b/tests/firewall/firewall_manager_test.py
@@ -53,7 +53,7 @@ def test_firewall_manager_update_existed():
     assert fm.host_controller.remove_rule.call_count == 0
 
 
-def test_firewall_manager_flush():
+def test_firewall_manager_cleanup():
     fm = SChainTestFirewallManager('test', 10000, 10064)
     rules = [
         SChainRule(10000, '2.2.2.2'),
@@ -63,6 +63,6 @@ def test_firewall_manager_flush():
     fm.add_rules(rules)
     fm.host_controller.add_rule(SChainRule(10072, '2.2.2.2'))
 
-    fm.flush()
+    fm.cleanup()
     assert list(fm.rules) == []
     assert fm.host_controller.has_rule(SChainRule(10072, '2.2.2.2'))
diff --git a/tests/firewall/nftables_test.py b/tests/firewall/nftables_test.py
index cee70eed..e77af1fa 100644
--- a/tests/firewall/nftables_test.py
+++ b/tests/firewall/nftables_test.py
@@ -7,6 +7,7 @@
 
 from core.schains.firewall.nftables import NFTablesController, NFT_CHAIN_BASE_PATH
 from core.schains.firewall.types import SChainRule
+from core.schains.firewall.utils import cleanup_firewall_for_schain
 from tools.helper import run_cmd
 
 
@@ -87,6 +88,9 @@ def test_create_delete_chain(filter_table, nft_chain_folder):
     manager.cleanup()
     chains = run_cmd(['nft', 'list', 'chains']).stdout.decode('utf-8')
     assert chains == 'table inet firewall {\n}\n'
+    assert os.path.isfile(nft_chain_path)
+
+    manager.remove_saved_rules()
     assert not os.path.isfile(nft_chain_path)
 
 
@@ -106,8 +110,21 @@ def test_saved_rules(filter_table, nft_chain_folder):
     assert not os.path.isfile(nft_chain_path)
 
 
+def test_cleanup_firewall_for_schain(filter_table, nft_chain_folder):
+    chain_name = 'test-chain'
+    nft_chain_path = os.path.join(NFT_CHAIN_BASE_PATH, f'skale-{chain_name}.conf')
+
+    manager = NFTablesController(chain=chain_name)
+    manager.create_chain(first_port=10000, last_port=10063)
+
+    cleanup_firewall_for_schain(schain_name=chain_name)
+    chains = run_cmd(['nft', 'list', 'chains']).stdout.decode('utf-8')
+    assert chains == 'table inet firewall {\n}\n'
+    assert not os.path.isfile(nft_chain_path)
+
+
 def add_remove_rule(srule, refresh):
-    manager = NFTablesController()
+    manager = NFTablesController(chain='test')
     manager.add_rule(srule)
     time.sleep(1)
     if not manager.has_rule(srule):
diff --git a/tests/schains/cleaner_test.py b/tests/schains/cleaner_test.py
index d16b41fd..ea45b25b 100644
--- a/tests/schains/cleaner_test.py
+++ b/tests/schains/cleaner_test.py
@@ -239,7 +239,8 @@ def test_get_schains_on_node(schain_dirs_for_monitor,
     ]).issubset(set(result))
 
 
-def test_remove_schain(skale, schain_db, node_config, dutils):
+@mock.patch('core.schains.cleaner.cleanup_firewall_for_schain')
+def test_remove_schain(cleanup_firewall_for_schain, skale, schain_db, node_config, dutils):
     schain_name = schain_db
     remove_schain(skale, node_config.id, schain_name, msg='Test remove_schain', dutils=dutils)
     container_name = SCHAIN_CONTAINER_NAME_TEMPLATE.format(schain_name)
@@ -250,7 +251,9 @@ def test_remove_schain(skale, schain_db, node_config, dutils):
     assert record.is_deleted is True
 
 
+@mock.patch('core.schains.cleaner.cleanup_firewall_for_schain')
 def test_cleanup_schain(
+    cleanup_firewall_rules,
     schain_db,
     node_config,
     schain_on_contracts,
diff --git a/tests/utils.py b/tests/utils.py
index 06b7e515..ac1f7f7d 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -233,6 +233,9 @@ class SChainTestFirewallManager(SChainFirewallManager):
     def create_host_controller(self):
         return HostTestFirewallController()
 
+    def cleanup(self):
+        self.remove_rules(self.rules)
+
 
 class SChainTestRuleController(SChainRuleController):
     def create_firewall_manager(self):
@@ -248,6 +251,9 @@ def is_persistent(self) -> bool:
     def is_inited(self) -> bool:
         return True
 
+    def cleanup(self) -> None:
+        self.firewall_manager.cleanup()
+
 
 def get_test_rule_controller(
     name,

From ad35813e2125263fbb081192bf388cb48976378a Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Wed, 22 Jan 2025 17:10:22 +0000
Subject: [PATCH 10/11] Improve rule controller test

---
 core/schains/checks.py                         | 4 ++--
 core/schains/firewall/firewall_manager.py      | 2 ++
 tests/firewall/default_rule_controller_test.py | 7 +++++++
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/core/schains/checks.py b/core/schains/checks.py
index 0a42bbdb..e3b95799 100644
--- a/core/schains/checks.py
+++ b/core/schains/checks.py
@@ -316,11 +316,11 @@ def firewall_rules(self) -> CheckRes:
                 base_port=base_port, own_ip=own_ip, node_ips=node_ips, sync_ip_ranges=ranges
             )
             logger.debug(f'Rule controller {self.rc.expected_rules()}')
-            data = {
+            data.update({
                 'inited': self.rc.is_inited(),
                 'rules': self.rc.is_rules_synced(),
                 'persistent': self.rc.is_persistent(),
-            }
+            })
             logger.debug('Firewall rules check: %s', data)
             status = all(data.values())
             return CheckRes(status=status, data=data)
diff --git a/core/schains/firewall/firewall_manager.py b/core/schains/firewall/firewall_manager.py
index f9d1bad2..5ae6cbe2 100644
--- a/core/schains/firewall/firewall_manager.py
+++ b/core/schains/firewall/firewall_manager.py
@@ -105,6 +105,8 @@ def create_host_controller(self) -> NFTablesController:
 
     def rules_saved(self) -> bool:
         saved = self.host_controller.get_saved_rules()
+        if saved == '':
+            return False
         return saved == self.host_controller.get_plain_chain_rules()
 
     def base_config_applied(self) -> bool:
diff --git a/tests/firewall/default_rule_controller_test.py b/tests/firewall/default_rule_controller_test.py
index cadd5f5d..ea211eab 100644
--- a/tests/firewall/default_rule_controller_test.py
+++ b/tests/firewall/default_rule_controller_test.py
@@ -37,6 +37,9 @@ def test_get_default_rule_controller(nft_chain_folder):
         node_ips,
         sync_ip_range
     )
+
+    assert rc.is_inited()
+    assert rc.is_persistent()
     assert rc.actual_rules() == []
     rc.sync()
     assert rc.expected_rules() == rc.actual_rules()
@@ -51,6 +54,10 @@ def test_get_default_rule_controller(nft_chain_folder):
     assert hm.add_rule.call_count == 0
     assert hm.remove_rule.call_count == 0
 
+    rc.cleanup()
+    assert not rc.is_inited()
+    assert not rc.is_persistent()
+
 
 def sync_rules(*args):
     rc = get_default_rule_controller(*args)

From 4e6fd905cc60d5e031e4ee27f7779b541ad17c2f Mon Sep 17 00:00:00 2001
From: badrogger <ivan@skalelabs.com>
Date: Wed, 22 Jan 2025 19:10:55 +0000
Subject: [PATCH 11/11] Fix tests

---
 core/schains/checks.py       | 2 +-
 tests/schains/checks_test.py | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/core/schains/checks.py b/core/schains/checks.py
index e3b95799..c0b5079f 100644
--- a/core/schains/checks.py
+++ b/core/schains/checks.py
@@ -304,7 +304,7 @@ def firewall_rules(self) -> CheckRes:
         data = {
             'inited': False,
             'rules': False,
-            'persistant': False,
+            'persistent': False,
         }
         if self.config:
             conf = self.cfm.skaled_config
diff --git a/tests/schains/checks_test.py b/tests/schains/checks_test.py
index f0d67f32..4ac14836 100644
--- a/tests/schains/checks_test.py
+++ b/tests/schains/checks_test.py
@@ -180,6 +180,8 @@ def test_volume_check(schain_checks, sample_false_checks, dutils):
 
 def test_firewall_rules_check(schain_checks, rules_unsynced_checks):
     schain_checks.rc.sync()
+    res = schain_checks.firewall_rules
+    print(res.data)
     assert schain_checks.firewall_rules
     assert not rules_unsynced_checks.firewall_rules.status