diff --git a/qubesagent/firewall.py b/qubesagent/firewall.py index 75b6292..b8af899 100755 --- a/qubesagent/firewall.py +++ b/qubesagent/firewall.py @@ -81,10 +81,8 @@ class FirewallWorker(object): raise NotImplementedError def get_connected_ips(self, family): - if family == 6: - return self.qdb.read('/connected-ips6').split() - else: - return self.qdb.read('/connected-ips').split() + ips = self.qdb.read('/connected-ips6' if family == 6 else '/connected-ips') + return ips.decode().split() def run_firewall_dir(self): """Run scripts dir contents, before user script""" diff --git a/qubesagent/test_firewall.py b/qubesagent/test_firewall.py index 0fa13c4..b3db626 100644 --- a/qubesagent/test_firewall.py +++ b/qubesagent/test_firewall.py @@ -332,10 +332,9 @@ class TestIptablesWorker(TestCase): ]) def test_008_update_connected_ips(self): - with patch.object(self.obj, 'get_connected_ips') as get_connected_ips: - get_connected_ips.return_value = ['10.137.0.1', '10.137.0.2'] - self.obj.called_commands[4] = [] - self.obj.update_connected_ips(4) + self.obj.qdb.entries['/connected-ips'] = b'10.137.0.1 10.137.0.2' + self.obj.called_commands[4] = [] + self.obj.update_connected_ips(4) self.assertEqual(self.obj.called_commands[4], [ ['-t', 'raw', '-F', 'QBS-PREROUTING'], @@ -529,10 +528,9 @@ class TestNftablesWorker(TestCase): ]) def test_008_update_connected_ips(self): - with patch.object(self.obj, 'get_connected_ips') as get_connected_ips: - get_connected_ips.return_value = ['10.137.0.1', '10.137.0.2'] - self.obj.loaded_rules = [] - self.obj.update_connected_ips(4) + self.obj.qdb.entries['/connected-ips'] = b'10.137.0.1 10.137.0.2' + self.obj.loaded_rules = [] + self.obj.update_connected_ips(4) self.assertEqual(self.obj.loaded_rules, [ 'flush chain ip qubes-firewall prerouting\n'