Sfoglia il codice sorgente

update_connected_ips: correctly handle byte-string

Pawel Marczewski 4 anni fa
parent
commit
e43fd2fc5a
2 ha cambiato i file con 8 aggiunte e 12 eliminazioni
  1. 2 4
      qubesagent/firewall.py
  2. 6 8
      qubesagent/test_firewall.py

+ 2 - 4
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"""

+ 6 - 8
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'