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
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"""

View File

@ -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'