SDK - Replaced insecure yaml.load with yaml.safe_load

This improves security and gets rid of security warnings.
See https://github.com/yaml/pyyaml/wiki/PyYAML-yaml.load(input)-Deprecation

components/ibm-components/ffdl/train/src/train.py

@@ -78,7 +78,7 @@
     ''' Update FfDL manifest with the corresponding object storage credentials '''
 
     f = open('manifest.yml', 'r')
-    manifest = yaml.load(f.read())
+    manifest = yaml.safe_load(f.read())
     f.close()
 
     manifest['data_stores'][0]['connection']['auth_url'] = s3_url

components/kubeflow/katib-launcher/src/launch_study_job.py

@@ -35,7 +35,7 @@ def yamlOrJsonStr(str):
     try:
         return json.loads(str)
     except:
-        return yaml.load(str)
+        return yaml.safe_load(str)
 
 def strToList(str):
     return str.split(",")
@@ -50,7 +50,7 @@ def _generate_studyjob_yaml(src_filename, name, namespace, optimizationtype, obj
                             metricsnames, parameterconfigs, nasConfig, workertemplatepath, mcollectortemplatepath, suggestionspec):
   """_generate_studyjob_yaml generates studyjob yaml file based on hp.template.yaml"""
   with open(src_filename, 'r') as f:
-    content = yaml.load(f)
+    content = yaml.safe_load(f)
 
   content['metadata']['name'] = name
   content['metadata']['namespace'] = namespace

components/kubeflow/launcher/src/launch_tf_job.py

@@ -49,7 +49,7 @@
 def _generate_train_yaml(src_filename, tfjob_ns, workers, pss, trainer_image, command):
   """_generate_train_yaml  generates train yaml files based on train.template.yaml"""
   with open(src_filename, 'r') as f:
-    content = yaml.load(f)
+    content = yaml.safe_load(f)
 
   content['metadata']['generateName'] = 'trainer-'
   content['metadata']['namespace'] = tfjob_ns

components/kubeflow/launcher/test/test_launcher.py

@@ -38,7 +38,7 @@ def test_yaml_generation_basic(self):
     args_list.append('--learning-rate=0.1')
     generated_yaml = train._generate_train_yaml(train_template_file, tfjob_ns, worker, pss, args_list)
     with open(os.path.join(test_data_dir, 'train_basic.yaml'), 'r') as f:
-      golden = yaml.load(f)
+      golden = yaml.safe_load(f)
     self.assertEqual(golden, generated_yaml)
 
   def test_yaml_generation_advanced(self):
@@ -53,7 +53,7 @@ def test_yaml_generation_advanced(self):
     args_list.append('--learning-rate=0.1')
     generated_yaml = train._generate_train_yaml(train_template_file, tfjob_ns, worker, pss, args_list)
     with open(os.path.join(test_data_dir, 'train_zero_worker.yaml'), 'r') as f:
-      golden = yaml.load(f)
+      golden = yaml.safe_load(f)
     self.assertEqual(golden, generated_yaml)
 
 if __name__ == '__main__':

sdk/python/kfp/_client.py

@@ -181,7 +181,7 @@ def _extract_pipeline_yaml(self, package_file):
           raise ValueError('Invalid package. Multiple yaml files in the package.')
 
         with tar.extractfile(all_yaml_files[0]) as f:
-          return yaml.load(f)
+          return yaml.safe_load(f)
     elif package_file.endswith('.zip'):
       with zipfile.ZipFile(package_file, 'r') as zip:
         all_yaml_files = [m for m in zip.namelist() if
@@ -193,10 +193,10 @@ def _extract_pipeline_yaml(self, package_file):
           raise ValueError('Invalid package. Multiple yaml files in the package.')
 
         with zip.open(all_yaml_files[0]) as f:
-          return yaml.load(f)
+          return yaml.safe_load(f)
     elif package_file.endswith('.yaml') or package_file.endswith('.yml'):
       with open(package_file, 'r') as f:
-        return yaml.load(f)
+        return yaml.safe_load(f)
     else:
       raise ValueError('The package_file '+ package_file + ' should ends with one of the following formats: [.tar.gz, .tgz, .zip, .yaml, .yml]')
 

sdk/python/kfp/components/_yaml_utils.py

@@ -18,7 +18,7 @@ def load_yaml(stream):
     #!!! Yaml should only be loaded using this function. Otherwise the dict ordering may be broken in Python versions prior to 3.6
     #See https://stackoverflow.com/questions/5121931/in-python-how-can-you-load-yaml-mappings-as-ordereddicts/21912744#21912744
 
-    def ordered_load(stream, Loader=yaml.Loader, object_pairs_hook=OrderedDict):
+    def ordered_load(stream, Loader=yaml.SafeLoader, object_pairs_hook=OrderedDict):
         class OrderedLoader(Loader):
             pass
         def construct_mapping(loader, node):

sdk/python/tests/compiler/compiler_tests.py

@@ -122,11 +122,11 @@ def test_operator_to_template(self):
   def _get_yaml_from_zip(self, zip_file):
     with zipfile.ZipFile(zip_file, 'r') as zip:
       with open(zip.extract(zip.namelist()[0]), 'r') as yaml_file:
-        return yaml.load(yaml_file)
+        return yaml.safe_load(yaml_file)
 
   def _get_yaml_from_tar(self, tar_file):
     with tarfile.open(tar_file, 'r:gz') as tar:
-      return yaml.load(tar.extractfile(tar.getmembers()[0]))
+      return yaml.safe_load(tar.extractfile(tar.getmembers()[0]))
 
   def test_basic_workflow(self):
     """Test compiling a basic workflow."""
@@ -139,7 +139,7 @@ def test_basic_workflow(self):
     try:
       compiler.Compiler().compile(basic.save_most_frequent_word, package_path)
       with open(os.path.join(test_data_dir, 'basic.yaml'), 'r') as f:
-        golden = yaml.load(f)
+        golden = yaml.safe_load(f)
       compiled = self._get_yaml_from_zip(package_path)
 
       self.maxDiff = None
@@ -166,7 +166,7 @@ def test_composing_workflow(self):
       compose_package_path = os.path.join(tmpdir, 'compose.zip')
       compiler.Compiler().compile(compose.download_save_most_frequent_word, compose_package_path)
       with open(os.path.join(test_data_dir, 'compose.yaml'), 'r') as f:
-        golden = yaml.load(f)
+        golden = yaml.safe_load(f)
       compiled = self._get_yaml_from_zip(compose_package_path)
 
       self.maxDiff = None
@@ -193,7 +193,7 @@ def test_package_compile(self):
           'dsl-compile', '--package', package_path, '--namespace', 'mypipeline',
           '--output', target_zip, '--function', 'download_save_most_frequent_word'])
       with open(os.path.join(test_data_dir, 'compose.yaml'), 'r') as f:
-        golden = yaml.load(f)
+        golden = yaml.safe_load(f)
       compiled = self._get_yaml_from_zip(target_zip)
 
       self.maxDiff = None
@@ -211,7 +211,7 @@ def _test_py_compile_zip(self, file_base_name):
       subprocess.check_call([
           'dsl-compile', '--py', py_file, '--output', target_zip])
       with open(os.path.join(test_data_dir, file_base_name + '.yaml'), 'r') as f:
-        golden = yaml.load(f)
+        golden = yaml.safe_load(f)
       compiled = self._get_yaml_from_zip(target_zip)
 
       self.maxDiff = None
@@ -228,7 +228,7 @@ def _test_py_compile_targz(self, file_base_name):
       subprocess.check_call([
           'dsl-compile', '--py', py_file, '--output', target_tar])
       with open(os.path.join(test_data_dir, file_base_name + '.yaml'), 'r') as f:
-        golden = yaml.load(f)
+        golden = yaml.safe_load(f)
       compiled = self._get_yaml_from_tar(target_tar)
       self.maxDiff = None
       self.assertEqual(golden, compiled)
@@ -244,10 +244,10 @@ def _test_py_compile_yaml(self, file_base_name):
       subprocess.check_call([
           'dsl-compile', '--py', py_file, '--output', target_yaml])
       with open(os.path.join(test_data_dir, file_base_name + '.yaml'), 'r') as f:
-        golden = yaml.load(f)
+        golden = yaml.safe_load(f)
 
       with open(os.path.join(test_data_dir, target_yaml), 'r') as f:
-        compiled = yaml.load(f)
+        compiled = yaml.safe_load(f)
 
       self.maxDiff = None
       self.assertEqual(golden, compiled)

sdk/python/tests/compiler/component_builder_test.py

@@ -313,7 +313,7 @@ def test_generate_kaniko_yaml(self):
     generated_yaml = builder._generate_kaniko_spec(namespace='default', arc_dockerfile_name='dockerfile',
                                                    gcs_path='gs://mlpipeline/kaniko_build.tar.gz', target_image='gcr.io/mlpipeline/kaniko_image:latest')
     with open(os.path.join(test_data_dir, 'kaniko.basic.yaml'), 'r') as f:
-      golden = yaml.load(f)
+      golden = yaml.safe_load(f)
 
     self.assertEqual(golden, generated_yaml)