Use Loader=SafeLoader for each yaml.load() call

Closes issue nitmir#15
fjalvarez09 · Jan 8, 2022 · 6c155a6 · 6c155a6
1 parent 79eb4e1
commit 6c155a6
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git a/policyd_rate_limit/tests/test_daemon.py b/policyd_rate_limit/tests/test_daemon.py
@@ -41,12 +41,12 @@ def test_main_unix_socket(self):
             self.base_test(cfg)
 
     def test_main_afinet_socket(self):
-        self.base_config["SOCKET"] = ("127.0.0.1", 27184)
+        self.base_config["SOCKET"] = ["127.0.0.1", 27184]
         with test_utils.lauch(self.base_config) as cfg:
             self.base_test(cfg)
 
     def test_main_afinet6_socket(self):
-        self.base_config["SOCKET"] = ("::1", 27184)
+        self.base_config["SOCKET"] = ["::1", 27184]
         with test_utils.lauch(self.base_config) as cfg:
             self.base_test(cfg)
 
@@ -174,10 +174,10 @@ def test_already_running(self):
                 pass
 
     def test_bad_socket_bind_address(self):
-        self.base_config["SOCKET"] = ("toto", 1234)
+        self.base_config["SOCKET"] = ["toto", 1234]
         with test_utils.lauch(self.base_config, get_process=True, no_wait=True) as p:
             self.assertEqual(p.wait(), 4)
-        self.base_config["SOCKET"] = ("192.168::1", 1234)
+        self.base_config["SOCKET"] = ["192.168::1", 1234]
         with test_utils.lauch(self.base_config, get_process=True, no_wait=True) as p:
             self.assertEqual(p.wait(), 6)
 

diff --git a/policyd_rate_limit/tests/utils.py b/policyd_rate_limit/tests/utils.py
@@ -58,6 +58,8 @@ def postfix_request(sasl_username="", client_address="127.0.0.1", protocol_state
 
 @contextmanager
 def sock(addr):
+    if isinstance(addr, list):
+        addr = tuple(addr)
     if isinstance(addr, str):
         s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
     elif '.' in addr[0]:
@@ -92,7 +94,7 @@ def gen_config(new_config):
         os.path.join(os.path.dirname(__file__), '..', 'policyd-rate-limit.yaml')
     )
     with open(default_config) as f:
-        config = yaml.load(f)
+        config = yaml.load(f, Loader=yaml.SafeLoader)
     config.update(new_config)
     cfg_path = tempfile.mktemp('.yaml')
     with open(cfg_path, 'w') as f:
@@ -156,7 +158,7 @@ def lauch(new_config, get_process=False, options=None, no_coverage=False, no_wai
     try:
         if cfg_path:
             with open(cfg_path) as f:
-                cfg = yaml.load(f)
+                cfg = yaml.load(f, Loader=yaml.SafeLoader)
             if not no_wait:
                 time.sleep(0.01)
                 for i in range(100):

diff --git a/policyd_rate_limit/utils.py b/policyd_rate_limit/utils.py
@@ -85,7 +85,7 @@ def __init__(self, config_file=None):
                     # new config file use yaml
                     else:
                         with open(config_file) as f:
-                            self._config = yaml.load(f)
+                            self._config = yaml.load(f, Loader=yaml.SafeLoader)
                     self.config_file = config_file
                     break
                 except PermissionError: