update_connected_ips: correctly handle byte-string

This commit is contained in:
Pawel Marczewski 2020-01-14 10:14:00 +01:00
parent 39885a4329
commit e43fd2fc5a
No known key found for this signature in database
GPG Key ID: DE42EE9B14F96465
2 changed files with 8 additions and 12 deletions

View File

@ -81,10 +81,8 @@ class FirewallWorker(object):
raise NotImplementedError raise NotImplementedError
def get_connected_ips(self, family): def get_connected_ips(self, family):
if family == 6: ips = self.qdb.read('/connected-ips6' if family == 6 else '/connected-ips')
return self.qdb.read('/connected-ips6').split() return ips.decode().split()
else:
return self.qdb.read('/connected-ips').split()
def run_firewall_dir(self): def run_firewall_dir(self):
"""Run scripts dir contents, before user script""" """Run scripts dir contents, before user script"""

View File

@ -332,10 +332,9 @@ class TestIptablesWorker(TestCase):
]) ])
def test_008_update_connected_ips(self): def test_008_update_connected_ips(self):
with patch.object(self.obj, 'get_connected_ips') as get_connected_ips: self.obj.qdb.entries['/connected-ips'] = b'10.137.0.1 10.137.0.2'
get_connected_ips.return_value = ['10.137.0.1', '10.137.0.2'] self.obj.called_commands[4] = []
self.obj.called_commands[4] = [] self.obj.update_connected_ips(4)
self.obj.update_connected_ips(4)
self.assertEqual(self.obj.called_commands[4], [ self.assertEqual(self.obj.called_commands[4], [
['-t', 'raw', '-F', 'QBS-PREROUTING'], ['-t', 'raw', '-F', 'QBS-PREROUTING'],
@ -529,10 +528,9 @@ class TestNftablesWorker(TestCase):
]) ])
def test_008_update_connected_ips(self): def test_008_update_connected_ips(self):
with patch.object(self.obj, 'get_connected_ips') as get_connected_ips: self.obj.qdb.entries['/connected-ips'] = b'10.137.0.1 10.137.0.2'
get_connected_ips.return_value = ['10.137.0.1', '10.137.0.2'] self.obj.loaded_rules = []
self.obj.loaded_rules = [] self.obj.update_connected_ips(4)
self.obj.update_connected_ips(4)
self.assertEqual(self.obj.loaded_rules, [ self.assertEqual(self.obj.loaded_rules, [
'flush chain ip qubes-firewall prerouting\n' 'flush chain ip qubes-firewall prerouting\n'