Fix code safety check (#1105)

* fix

* linting

* fix

* fix

* linting

Author: SamitHuang

examples/hunyuan_dit/hydit/data_loader/csv2arrow.py

@@ -20,7 +20,7 @@ def parse_data(data):
 
         with open(img_path, "rb") as fp:
             image = fp.read()
-            md5 = hashlib.md5(image).hexdigest()
+            md5 = hashlib.md5(image, usedforsecurity=False).hexdigest()
 
         with Image.open(img_path) as f:
             width, height = f.size

examples/opensora_pku/opensora/models/causalvideovae/eval/fvd/styleganv/fvd.py

@@ -2,6 +2,8 @@
 # https://github.com/PKU-YuanGroup/Open-Sora-Plan/blob/main/opensora/models/causalvideovae/eval/fvd/styleganv/fvd.py
 import math
 import os
+import shlex
+import subprocess
 
 from mindspore import context, export, load, mint, nn, ops
 
@@ -23,7 +25,7 @@ def load_i3d_pretrained(bs=1):
     if not os.path.exists(mindir_filepath):
         if not os.path.exists(filepath):
             print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.")
-            os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}")
+            subprocess.run(shlex.split(f"wget {i3D_WEIGHTS_URL} -O {filepath}"), shell=False)
         if not os.path.exists(onnx_filepath):
             # convert torch jit model to onnx model
             model = torch.jit.load(filepath).eval()

examples/opensora_pku/opensora/models/causalvideovae/eval/fvd/videogpt/fvd.py

@@ -3,6 +3,8 @@
 
 import math
 import os
+import shlex
+import subprocess
 
 import numpy as np
 
@@ -26,7 +28,7 @@ def load_i3d_pretrained():
     if not os.path.exists(ms_filepath):
         if not os.path.exists(filepath):
             print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.")
-            os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}")
+            subprocess.run(shlex.split(f"wget {i3D_WEIGHTS_URL} -O {filepath}"), shell=False)
         # convert torch ckpt to mindspore ckpt
         state_dict = torch.load_state_dict(torch.load(filepath))
         raise ValueError("Not converted")

examples/stable_diffusion_v2/ldm/data/sync_data.py

@@ -32,15 +32,42 @@ def get_rank_id():
     return int(global_rank_id)
 
 
+def is_safe_member(member, target_dir):
+    member_path = os.path.join(target_dir, member.name)
+    abs_target_dir = os.path.abspath(target_dir)
+    abs_member_path = os.path.abspath(member_path)
+
+    if not abs_member_path.startswith(abs_target_dir):
+        return False
+
+    if member.name.startswith("/") or ".." in member.name:
+        return False
+
+    if member.islnk() or member.issym():
+        return False
+
+    return True
+
+
+def safe_members(tar, target_dir):
+    for member in tar.getmembers():
+        if is_safe_member(member, target_dir):
+            yield member
+        else:
+            print(f"Discarding unsafe member: {member.name}")
+
+
 def extract_tar(file_path):
     try:
         with tarfile.open(file_path, "r") as archive:
             if "/" not in archive.getnames()[1]:
                 subfolder_path = file_path[:-4]
                 os.makedirs(subfolder_path, exist_ok=True)
-                archive.extractall(subfolder_path)
+                archive.extractall(subfolder_path, members=safe_members(archive, subfolder_path))
             else:
-                archive.extractall(os.path.dirname(file_path))
+                archive.extractall(
+                    os.path.dirname(file_path), members=safe_members(archive, os.path.dirname(file_path))
+                )
         os.remove(file_path)
         _logger.info("finish extract: " + file_path)
         return True

examples/stable_diffusion_v2/tests/st/test_train_infer_dummy.py

@@ -1,4 +1,5 @@
 import os
+import shlex
 import subprocess
 import sys
 
@@ -19,7 +20,7 @@ def test_train_infer(use_ema, finetuning):
     epochs = 1
     image_size = 512
     if os.path.exists(output_path):
-        os.system(f"rm {output_path} -rf")
+        subprocess.run(shlex.split(f"rm {output_path} -rf"), shell=False)
     if finetuning == "Vanilla":
         use_lora = False
     elif finetuning == "LoRA":
@@ -65,7 +66,7 @@ def test_train_infer_DreamBooth(use_ema):
     epochs = 1
     image_size = 512
     if os.path.exists(output_path):
-        os.system(f"rm {output_path} -rf")
+        subprocess.run(shlex.split(f"rm {output_path} -rf"), shell=False)
     cmd = (
         f"python train_dreambooth.py --mode=0 --instance_data_dir={data_path} --instance_prompt='{instance_prompt}' "
         f"--model_config={model_config_file} --class_data_dir={data_path} --class_prompt='{instance_prompt}' "

examples/stable_diffusion_v2/tools/eval/fid/utils.py

@@ -6,7 +6,9 @@
 import os
 import pathlib
 import ssl
+import subprocess
 import tarfile
+import time
 import urllib
 import urllib.error
 import urllib.request
@@ -85,6 +87,39 @@ def detect_file_type(filename: str):  # pylint: disable=inconsistent-return-stat
         return suffix, None, suffix
 
 
+def download_weights(url, dest):
+    start = time.time()
+    print("downloading url: ", url)
+    print("downloading to: ", dest)
+    subprocess.check_call(["pget", "-x", url, dest], close_fds=False)
+    print("downloading took: ", time.time() - start)
+
+
+def is_safe_member(member, target_dir):
+    member_path = os.path.join(target_dir, member.name)
+    abs_target_dir = os.path.abspath(target_dir)
+    abs_member_path = os.path.abspath(member_path)
+
+    if not abs_member_path.startswith(abs_target_dir):
+        return False
+
+    if member.name.startswith("/") or ".." in member.name:
+        return False
+
+    if member.islnk() or member.issym():
+        return False
+
+    return True
+
+
+def safe_members(tar, target_dir):
+    for member in tar.getmembers():
+        if is_safe_member(member, target_dir):
+            yield member
+        else:
+            print(f"Discarding unsafe member: {member.name}")
+
+
 class Download:
     """Base utility class for downloading."""
 
@@ -96,7 +131,7 @@ class Download:
     @staticmethod
     def calculate_md5(file_path: str, chunk_size: int = 1024 * 1024) -> str:
         """Calculate md5 value."""
-        md5 = hashlib.md5()
+        md5 = hashlib.md5(usedforsecurity=False)
         with open(file_path, "rb") as fp:
             for chunk in iter(lambda: fp.read(chunk_size), b""):
                 md5.update(chunk)
@@ -111,15 +146,15 @@ def extract_tar(from_path: str, to_path: Optional[str] = None, compression: Opti
         """Extract tar format file."""
 
         with tarfile.open(from_path, f"r:{compression[1:]}" if compression else "r") as tar:
-            tar.extractall(to_path)
+            tar.extract_all(tar, members=safe_members(tar, to_path))
 
     @staticmethod
     def extract_zip(from_path: str, to_path: Optional[str] = None, compression: Optional[str] = None) -> None:
         """Extract zip format file."""
 
         compression_mode = zipfile.ZIP_BZIP2 if compression else zipfile.ZIP_STORED
         with zipfile.ZipFile(from_path, "r", compression=compression_mode) as zip_file:
-            zip_file.extractall(to_path)
+            zipfile.extract_all(zip_file, members=safe_members(zip_file, to_path))
 
     def extract_archive(self, from_path: str, to_path: str = None) -> str:
         """Extract and  archive from path to path."""

examples/stable_diffusion_v2/utils/download.py

@@ -1,11 +1,14 @@
 """Utility of downloading"""
+
 import bz2
 import gzip
 import hashlib
 import logging
 import os
 import ssl
+import subprocess
 import tarfile
+import time
 import urllib
 import urllib.error
 import urllib.request
@@ -32,6 +35,39 @@ def set_default_download_root(path):
     _DEFAULT_DOWNLOAD_ROOT = path
 
 
+def download_weights(url, dest):
+    start = time.time()
+    print("downloading url: ", url)
+    print("downloading to: ", dest)
+    subprocess.check_call(["pget", "-x", url, dest], close_fds=False)
+    print("downloading took: ", time.time() - start)
+
+
+def is_safe_member(member, target_dir):
+    member_path = os.path.join(target_dir, member.name)
+    abs_target_dir = os.path.abspath(target_dir)
+    abs_member_path = os.path.abspath(member_path)
+
+    if not abs_member_path.startswith(abs_target_dir):
+        return False
+
+    if member.name.startswith("/") or ".." in member.name:
+        return False
+
+    if member.islnk() or member.issym():
+        return False
+
+    return True
+
+
+def safe_members(tar, target_dir):
+    for member in tar.getmembers():
+        if is_safe_member(member, target_dir):
+            yield member
+        else:
+            print(f"Discarding unsafe member: {member.name}")
+
+
 class DownLoad:
     """Base utility class for downloading."""
 
@@ -43,7 +79,7 @@ class DownLoad:
     @staticmethod
     def calculate_md5(file_path: str, chunk_size: int = 1024 * 1024) -> str:
         """Calculate md5 value."""
-        md5 = hashlib.md5()
+        md5 = hashlib.md5(usedforsecurity=False)
         with open(file_path, "rb") as fp:
             for chunk in iter(lambda: fp.read(chunk_size), b""):
                 md5.update(chunk)
@@ -58,15 +94,15 @@ def extract_tar(from_path: str, to_path: Optional[str] = None, compression: Opti
         """Extract tar format file."""
 
         with tarfile.open(from_path, f"r:{compression[1:]}" if compression else "r") as tar:
-            tar.extractall(to_path)
+            tar.extract_all(tar, members=safe_members(tar, to_path))
 
     @staticmethod
     def extract_zip(from_path: str, to_path: Optional[str] = None, compression: Optional[str] = None) -> None:
         """Extract zip format file."""
 
         compression_mode = zipfile.ZIP_BZIP2 if compression else zipfile.ZIP_STORED
         with zipfile.ZipFile(from_path, "r", compression=compression_mode) as zip_file:
-            zip_file.extractall(to_path)
+            zipfile.extract_all(zip_file, members=safe_members(zip_file, to_path))
 
     def extract_archive(self, from_path: str, to_path: str = None) -> str:
         """Extract and  archive from path to path."""

examples/stable_diffusion_xl/tools/rank_table_generation/hccl_tools.py

@@ -16,7 +16,9 @@
 import json
 import os
 import re
+import shlex
 import socket
+import subprocess
 import sys
 from argparse import ArgumentParser
 from typing import Any, Dict
@@ -133,8 +135,11 @@ def main():
     device_ips: Dict[Any, Any] = {}
     try:
         for device_id in device_num_list:
-            ret = os.popen("hccn_tool -i %d -ip -g" % device_id).readlines()
-            device_ips[str(device_id)] = ret[0].split(":")[1].replace("\n", "")
+            cmd = "hccn_tool -i %d -ip -g" % device_id
+            result = subprocess.run(shlex.split(cmd), capture_output=True, text=True, shell=False)
+            output = result.stdout.strip()
+            device_ips[str(device_id)] = output.split(":")[1].replace("\n", "")
+
     except IndexError:
         print("Failed to call hccn_tool, try to read /etc/hccn.conf instead")
         try:

examples/t2v_turbo/utils/download.py

@@ -43,6 +43,31 @@ def download_weights(url, dest):
     print("downloading took: ", time.time() - start)
 
 
+def is_safe_member(member, target_dir):
+    member_path = os.path.join(target_dir, member.name)
+    abs_target_dir = os.path.abspath(target_dir)
+    abs_member_path = os.path.abspath(member_path)
+
+    if not abs_member_path.startswith(abs_target_dir):
+        return False
+
+    if member.name.startswith("/") or ".." in member.name:
+        return False
+
+    if member.islnk() or member.issym():
+        return False
+
+    return True
+
+
+def safe_members(tar, target_dir):
+    for member in tar.getmembers():
+        if is_safe_member(member, target_dir):
+            yield member
+        else:
+            print(f"Discarding unsafe member: {member.name}")
+
+
 class DownLoad:
     """Base utility class for downloading."""
 
@@ -54,7 +79,7 @@ class DownLoad:
     @staticmethod
     def calculate_md5(file_path: str, chunk_size: int = 1024 * 1024) -> str:
         """Calculate md5 value."""
-        md5 = hashlib.md5()
+        md5 = hashlib.md5(usedforsecurity=False)
         with open(file_path, "rb") as fp:
             for chunk in iter(lambda: fp.read(chunk_size), b""):
                 md5.update(chunk)
@@ -69,15 +94,15 @@ def extract_tar(from_path: str, to_path: Optional[str] = None, compression: Opti
         """Extract tar format file."""
 
         with tarfile.open(from_path, f"r:{compression[1:]}" if compression else "r") as tar:
-            tar.extractall(to_path)
+            tar.extract_all(tar, members=safe_members(tar, to_path))
 
     @staticmethod
     def extract_zip(from_path: str, to_path: Optional[str] = None, compression: Optional[str] = None) -> None:
         """Extract zip format file."""
 
         compression_mode = zipfile.ZIP_BZIP2 if compression else zipfile.ZIP_STORED
         with zipfile.ZipFile(from_path, "r", compression=compression_mode) as zip_file:
-            zip_file.extractall(to_path)
+            zipfile.extract_all(zip_file, members=safe_members(zip_file, to_path))
 
     def extract_archive(self, from_path: str, to_path: str = None) -> str:
         """Extract and  archive from path to path."""

examples/venhancer/inference_utils.py

@@ -4,6 +4,8 @@
 
 import logging
 import os
+import shlex
+import shutil
 import subprocess
 import tempfile
 
@@ -109,8 +111,10 @@ def save_video(video, save_dir, file_name, fps=16.0):
     tmp_path = os.path.join(save_dir, "tmp.mp4")
     cmd = f"ffmpeg -y -f image2 -framerate {fps} -i {temp_dir}/%06d.png \
       -crf 17 -pix_fmt yuv420p {tmp_path}"
-    status, output = subprocess.getstatusoutput(cmd)
+    result = subprocess.run(shlex.split(cmd), capture_output=True, text=True, shell=False)
+    status = result.returncode
+    output = result.stdout
     if status != 0:
         logger.error(f"Save Video Error with {output}")
-    os.system(f"rm -rf {temp_dir}")
+    shutil.rmtree(temp_dir)
     os.rename(tmp_path, output_path)