Merge pull request #1 from Anon-Artist/RCE-fix

Rce fix
microsoft · Jan 17, 2021 · ea3beba · ea3beba
2 parents fbbe14d + a2d815d
commit ea3beba
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/nni/tools/nnictl/common_utils.py b/nni/tools/nnictl/common_utils.py
@@ -21,7 +21,7 @@ def get_yml_content(file_path):
     '''Load yaml file content'''
     try:
         with open(file_path, 'r') as file:
-            return yaml.load(file, Loader=yaml.Loader)
+            return yaml.load(file, Loader=yaml.SafeLoader)
     except yaml.scanner.ScannerError as err:
         print_error('yaml file format error!')
         print_error(err)

diff --git a/nni/tools/package_utils/__init__.py b/nni/tools/package_utils/__init__.py
@@ -216,7 +216,7 @@ def read_registerd_algo_meta():
     config_file = get_registered_algo_config_path()
     if os.path.exists(config_file):
         with open(config_file, 'r') as f:
-            config = yaml.load(f, Loader=yaml.Loader)
+            config = yaml.load(f, Loader=yaml.SafeLoader)
     else:
         config = defaultdict(list)
     for t in ALGO_TYPES: