From 3af2bfb846c81ed695a39ad06e8bde03442f1c80 Mon Sep 17 00:00:00 2001 From: badrogger Date: Fri, 3 Jan 2025 11:54:43 +0000 Subject: [PATCH 1/6] Modify flake8 config to trigger build --- .flake8 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.flake8 b/.flake8 index 12bd9cf3..a34d601c 100644 --- a/.flake8 +++ b/.flake8 @@ -1,3 +1,3 @@ [flake8] max-line-length = 100 -exclude = .git,__pycache__,docs/source/conf.py,old,build,dist,venv,helper-scripts +exclude = .git,__pycache__,docs/source/conf.py,old,build,dist,venv,helper-scripts,.venv From f33ec05ac617850732bd0da26236bcdff0c1a678 Mon Sep 17 00:00:00 2001 From: badrogger Date: Tue, 7 Jan 2025 19:45:35 +0000 Subject: [PATCH 2/6] Do not flush rules, remove the redundant --- node_cli/core/nftables.py | 102 +++++++++++++++++++++++++- node_cli/migrations/focal_to_jammy.py | 9 ++- 2 files changed, 108 insertions(+), 3 deletions(-) diff --git a/node_cli/core/nftables.py b/node_cli/core/nftables.py index de411ddb..38eaba5a 100644 --- a/node_cli/core/nftables.py +++ b/node_cli/core/nftables.py @@ -162,7 +162,7 @@ def rule_exists(self, chain: str, new_rule_expr: list[dict]) -> bool: return True return False - def add_drop_rule_if_node_exists(self, protocol: str) -> None: + def add_drop_rule_if_not_exists(self, protocol: str) -> None: expr = [ { "match": { @@ -197,6 +197,45 @@ def add_drop_rule_if_node_exists(self, protocol: str) -> None: self.execute_cmd(cmd) logger.info('Added drop rule for %s', protocol) + def remove_drop_rule_if_exists(self, protocol: str) -> None: + expr = [ + { + "match": { + "op": "==", + "left": { + "payload": { + "protocol": "ip", + "field": "protocol" + } + }, + "right": protocol + } + }, + {'counter': None}, + {"drop": None} + ] + + # Check if the drop rule exists before attempting to remove it + if self.rule_exists(self.chain, expr): + cmd = { + 'nftables': [ + { + 'delete': { + 'rule': { + 'family': self.family, + 'table': self.table, + 'chain': self.chain, + 'expr': expr, + } + } + } + ] + } + self.execute_cmd(cmd) + logger.info('Removed drop rule for %s', protocol) + else: + logger.info('Drop rule does not exist for %s', protocol) + def add_rule_if_not_exists(self, rule: Rule) -> None: expr = [] @@ -249,6 +288,56 @@ def add_rule_if_not_exists(self, rule: Rule) -> None: 'Rule already exists in chain %s: %s port %s', rule.chain, rule.protocol, rule.port ) + def remove_rule_if_exists(self, rule: Rule) -> None: + expr = [] + + if rule.protocol in ['tcp', 'udp']: + if rule.port: + expr.append( + { + 'match': { + 'left': {'payload': {'protocol': rule.protocol, 'field': 'dport'}}, + 'op': '==', + 'right': rule.port, + } + } + ) + elif rule.protocol == 'icmp' and rule.icmp_type: + expr.append( + { + 'match': { + 'left': {'payload': {'protocol': 'icmp', 'field': 'type'}}, + 'op': '==', + 'right': rule.icmp_type, + } + } + ) + + # Check if the rule exists before attempting to remove it + if self.rule_exists(rule.chain, expr): + cmd = { + 'nftables': [ + { + 'delete': { + 'rule': { + 'family': self.family, + 'table': self.table, + 'chain': rule.chain, + 'expr': expr, + } + } + } + ] + } + self.execute_cmd(cmd) + logger.info( + 'Removed rule from chain %s: %s port %s', rule.chain, rule.protocol, rule.port + ) + else: + logger.info( + 'Rule does not exist in chain %s: %s port %s', rule.chain, rule.protocol, rule.port + ) + def add_connection_tracking_rule(self, chain: str) -> None: expr = [ { @@ -354,13 +443,22 @@ def setup_firewall(self, enable_monitoring: bool = False) -> None: ) ) - self.add_drop_rule_if_node_exists(protocol='udp') + self.add_drop_rule_if_not_exists(protocol='udp') except Exception as e: logger.error('Failed to setup firewall: %s', e) raise NFTablesError(e) logger.info('Firewall rules are configured') + def cleanup_rules(self): + """ Cleanups all node-cli generated rules """ + self.remove_drop_rule_if_exists('tcp') + self.remove_drop_rule_if_exists('udp') + tcp_ports = [get_ssh_port(), 53, 443, 3009, 8080, 9100] + for port in tcp_ports: + self.remove_rule_if_exists(Rule(chain=self.chain, protocol='tcp', port=port)) + self.remove_rule_if_exists(Rule(chain=self.chain, protocol='udp', port=53)) + def flush_chain(self, chain: str) -> None: """Remove all rules from a specific chain""" json_cmd = { diff --git a/node_cli/migrations/focal_to_jammy.py b/node_cli/migrations/focal_to_jammy.py index eb63b6e2..6be24b49 100644 --- a/node_cli/migrations/focal_to_jammy.py +++ b/node_cli/migrations/focal_to_jammy.py @@ -118,7 +118,14 @@ def migrate() -> None: ssh_port = get_ssh_port() logger.info('Running migration from focal to jammy') remove_old_iptables_rules(ssh_port) + logger.info('Flushing nftables rules generated by release upgrade') nft = NFTablesManager(family='ip', table='filter') - nft.flush_chain(IPTABLES_CHAIN) + nft.cleanup_rules() + + # Logging rules after migration + res = run_cmd(['nftables', 'list', 'ruleset']) + plain_rules = res.stdout.decode('utf-8') + logger.debug(plain_rules) + logger.info('Migration from focal to jammy completed') From 26bfa2ceae7e10444e0e3240112d0018e83f45b7 Mon Sep 17 00:00:00 2001 From: badrogger Date: Wed, 8 Jan 2025 13:40:41 +0000 Subject: [PATCH 3/6] Fix tests --- node_cli/migrations/focal_to_jammy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/node_cli/migrations/focal_to_jammy.py b/node_cli/migrations/focal_to_jammy.py index 6be24b49..c63db869 100644 --- a/node_cli/migrations/focal_to_jammy.py +++ b/node_cli/migrations/focal_to_jammy.py @@ -124,7 +124,7 @@ def migrate() -> None: nft.cleanup_rules() # Logging rules after migration - res = run_cmd(['nftables', 'list', 'ruleset']) + res = run_cmd(['nft', 'list', 'ruleset']) plain_rules = res.stdout.decode('utf-8') logger.debug(plain_rules) From a32e71c1cfadb049095e225850ab59a1062ca628 Mon Sep 17 00:00:00 2001 From: badrogger Date: Wed, 8 Jan 2025 13:58:41 +0000 Subject: [PATCH 4/6] Improve tests --- tests/core/migration_test.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/core/migration_test.py b/tests/core/migration_test.py index c3f20d57..66759d5e 100644 --- a/tests/core/migration_test.py +++ b/tests/core/migration_test.py @@ -1,6 +1,6 @@ import pytest -from node_cli.migrations.focal_to_jammy import migrate +from node_cli.migrations.focal_to_jammy import migrate, NFTablesManager from node_cli.utils.helper import run_cmd @@ -41,3 +41,5 @@ def test_migration(base_rules): res = run_cmd(['iptables', '-S']) output = res.stdout.decode('utf-8') assert output == f'-P INPUT ACCEPT\n-P FORWARD ACCEPT\n-P OUTPUT ACCEPT\n-N {CUSTOM_CHAIN_NAME}\n-A {CUSTOM_CHAIN_NAME} -p tcp -m tcp --dport 2222 -j ACCEPT\n' # noqa + nft = NFTablesManager(family='ip', table='filter') + assert nft.get_rules(chain='INPUT') == [] From 6acc02183e1b9fd4f9bbe81f2a4d19d13ca35757 Mon Sep 17 00:00:00 2001 From: badrogger Date: Wed, 8 Jan 2025 19:09:39 +0000 Subject: [PATCH 5/6] Small improvements --- node_cli/core/nftables.py | 22 +++++++++++----------- tests/core/nftables_test.py | 6 +++--- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/node_cli/core/nftables.py b/node_cli/core/nftables.py index 38eaba5a..eb12374c 100644 --- a/node_cli/core/nftables.py +++ b/node_cli/core/nftables.py @@ -162,7 +162,7 @@ def rule_exists(self, chain: str, new_rule_expr: list[dict]) -> bool: return True return False - def add_drop_rule_if_not_exists(self, protocol: str) -> None: + def add_drop_rule(self, protocol: str) -> None: expr = [ { "match": { @@ -197,7 +197,7 @@ def add_drop_rule_if_not_exists(self, protocol: str) -> None: self.execute_cmd(cmd) logger.info('Added drop rule for %s', protocol) - def remove_drop_rule_if_exists(self, protocol: str) -> None: + def remove_drop_rule(self, protocol: str) -> None: expr = [ { "match": { @@ -236,7 +236,7 @@ def remove_drop_rule_if_exists(self, protocol: str) -> None: else: logger.info('Drop rule does not exist for %s', protocol) - def add_rule_if_not_exists(self, rule: Rule) -> None: + def add_rule(self, rule: Rule) -> None: expr = [] if rule.protocol in ['tcp', 'udp']: @@ -428,14 +428,14 @@ def setup_firewall(self, enable_monitoring: bool = False) -> None: if enable_monitoring: tcp_ports.extend([8080, 9100]) for port in tcp_ports: - self.add_rule_if_not_exists(Rule(chain=self.chain, protocol='tcp', port=port)) + self.add_rule(Rule(chain=self.chain, protocol='tcp', port=port)) - self.add_rule_if_not_exists(Rule(chain=self.chain, protocol='udp', port=53)) + self.add_rule(Rule(chain=self.chain, protocol='udp', port=53)) self.add_loopback_rule(chain=self.chain) icmp_types = ['destination-unreachable', 'source-quench', 'time-exceeded'] for icmp_type in icmp_types: - self.add_rule_if_not_exists( + self.add_rule( Rule( chain=self.chain, protocol='icmp', @@ -443,7 +443,7 @@ def setup_firewall(self, enable_monitoring: bool = False) -> None: ) ) - self.add_drop_rule_if_not_exists(protocol='udp') + self.add_drop_rule(protocol='udp') except Exception as e: logger.error('Failed to setup firewall: %s', e) @@ -452,12 +452,12 @@ def setup_firewall(self, enable_monitoring: bool = False) -> None: def cleanup_rules(self): """ Cleanups all node-cli generated rules """ - self.remove_drop_rule_if_exists('tcp') - self.remove_drop_rule_if_exists('udp') + self.remove_drop_rule('tcp') + self.remove_drop_rule('udp') tcp_ports = [get_ssh_port(), 53, 443, 3009, 8080, 9100] for port in tcp_ports: - self.remove_rule_if_exists(Rule(chain=self.chain, protocol='tcp', port=port)) - self.remove_rule_if_exists(Rule(chain=self.chain, protocol='udp', port=53)) + self.remove_rule(Rule(chain=self.chain, protocol='tcp', port=port)) + self.remove_rule(Rule(chain=self.chain, protocol='udp', port=53)) def flush_chain(self, chain: str) -> None: """Remove all rules from a specific chain""" diff --git a/tests/core/nftables_test.py b/tests/core/nftables_test.py index e63e8d81..28fc68f2 100644 --- a/tests/core/nftables_test.py +++ b/tests/core/nftables_test.py @@ -112,12 +112,12 @@ def test_create_chain_if_not_exists(mock_exists, mock_execute, nft_manager): ) @patch.object(NFTablesManager, 'execute_cmd') @patch.object(NFTablesManager, 'rule_exists') -def test_add_rule_if_not_exists(mock_exists, mock_execute, nft_manager, rule_data): +def test_add_rule(mock_exists, mock_execute, nft_manager, rule_data): """Test rule addition with different types""" mock_exists.return_value = False rule = Rule(**rule_data) - nft_manager.add_rule_if_not_exists(rule) + nft_manager.add_rule(rule) mock_execute.assert_called_once() @@ -138,4 +138,4 @@ def test_invalid_protocol(nft_manager): """Test adding rule with invalid protocol""" rule = Rule(chain='INPUT', protocol='invalid', port=80) with pytest.raises(Exception): - nft_manager.add_rule_if_not_exists(rule) + nft_manager.add_rule(rule) From 9129c585d621bedaabf9f087ee7a914cbc3c36f4 Mon Sep 17 00:00:00 2001 From: badrogger Date: Wed, 8 Jan 2025 19:40:07 +0000 Subject: [PATCH 6/6] Fix tests --- node_cli/core/nftables.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/node_cli/core/nftables.py b/node_cli/core/nftables.py index eb12374c..9b467963 100644 --- a/node_cli/core/nftables.py +++ b/node_cli/core/nftables.py @@ -288,7 +288,7 @@ def add_rule(self, rule: Rule) -> None: 'Rule already exists in chain %s: %s port %s', rule.chain, rule.protocol, rule.port ) - def remove_rule_if_exists(self, rule: Rule) -> None: + def remove_rule(self, rule: Rule) -> None: expr = [] if rule.protocol in ['tcp', 'udp']: