mindstudio-probe 1.0.4__py3-none-any.whl → 1.1.1__py3-none-any.whl

msprobe/core/common/file_utils.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2022-2023. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
@@ -13,15 +12,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
+
 import csv
 import fcntl
 import os
+import stat
 import json
 import re
 import shutil
+from datetime import datetime, timezone
+from dateutil import parser
+import OpenSSL
 import yaml
 import numpy as np
+import pandas as pd
 
 from msprobe.core.common.log import logger
 from msprobe.core.common.exceptions import FileCheckException
@@ -66,9 +70,11 @@ class FileChecker:
         self.check_path_ability()
         if self.is_script:
             check_path_owner_consistent(self.file_path)
-        check_path_pattern_vaild(self.file_path)
+        check_path_pattern_valid(self.file_path)
         check_common_file_size(self.file_path)
         check_file_suffix(self.file_path, self.file_type)
+        if self.path_type == FileCheckConst.FILE:
+            check_dirpath_before_read(self.file_path)
         return self.file_path
 
     def check_path_ability(self):
@@ -121,9 +127,10 @@ class FileOpen:
         self.file_path = os.path.realpath(self.file_path)
         check_path_length(self.file_path)
         self.check_ability_and_owner()
-        check_path_pattern_vaild(self.file_path)
+        check_path_pattern_valid(self.file_path)
         if os.path.exists(self.file_path):
             check_common_file_size(self.file_path)
+            check_dirpath_before_read(self.file_path)
 
     def check_ability_and_owner(self):
         if self.mode in self.SUPPORT_READ_MODE:
@@ -187,12 +194,12 @@ def check_other_user_writable(path):
 
 def check_path_owner_consistent(path):
     file_owner = os.stat(path).st_uid
-    if file_owner != os.getuid():
+    if file_owner != os.getuid() and os.getuid() != 0:
         logger.error('The file path %s may be insecure because is does not belong to you.' % path)
         raise FileCheckException(FileCheckException.FILE_PERMISSION_ERROR)
 
 
-def check_path_pattern_vaild(path):
+def check_path_pattern_valid(path):
     if not re.match(FileCheckConst.FILE_VALID_PATTERN, path):
         logger.error('The file path %s contains special characters.' % (path))
         raise FileCheckException(FileCheckException.ILLEGAL_PATH_ERROR)
@@ -214,7 +221,8 @@ def check_common_file_size(file_path):
         for suffix, max_size in FileCheckConst.FILE_SIZE_DICT.items():
             if file_path.endswith(suffix):
                 check_file_size(file_path, max_size)
-                break
+                return
+        check_file_size(file_path, FileCheckConst.COMMOM_FILE_SIZE)
 
 
 def check_file_suffix(file_path, file_suffix):
@@ -235,9 +243,18 @@ def check_path_type(file_path, file_type):
             raise FileCheckException(FileCheckException.INVALID_FILE_ERROR)
 
 
+def check_others_writable(directory):
+    dir_stat = os.stat(directory)
+    is_writable = (
+        bool(dir_stat.st_mode & stat.S_IWGRP) or  # 组可写
+        bool(dir_stat.st_mode & stat.S_IWOTH)     # 其他用户可写
+    )
+    return is_writable
+
+
 def make_dir(dir_path):
-    dir_path = os.path.realpath(dir_path)
     check_path_before_create(dir_path)
+    dir_path = os.path.realpath(dir_path)
     if os.path.isdir(dir_path):
         return
     try:
@@ -259,8 +276,9 @@ def create_directory(dir_path):
     Exception Description:
         when invalid data throw exception
     """
-    dir_path = os.path.realpath(dir_path)
+    check_link(dir_path)
     check_path_before_create(dir_path)
+    dir_path = os.path.realpath(dir_path)
     parent_dir = os.path.dirname(dir_path)
     if not os.path.isdir(parent_dir):
         create_directory(parent_dir)
@@ -268,6 +286,7 @@ def create_directory(dir_path):
 
 
 def check_path_before_create(path):
+    check_link(path)
     if path_len_exceeds_limit(path):
         raise FileCheckException(FileCheckException.ILLEGAL_PATH_ERROR, 'The file path length exceeds limit.')
 
@@ -276,6 +295,17 @@ def check_path_before_create(path):
                                  'The file path {} contains special characters.'.format(path))
 
 
+def check_dirpath_before_read(path):
+    path = os.path.realpath(path)
+    dirpath = os.path.dirname(path)
+    if check_others_writable(dirpath):
+        logger.warning(f"The directory is writable by others: {dirpath}.")
+    try:
+        check_path_owner_consistent(dirpath)
+    except FileCheckException:
+        logger.warning(f"The directory {dirpath} is not yours.")
+    
+
 def check_file_or_directory_path(path, isdir=False):
     """
     Function Description:
@@ -322,7 +352,7 @@ def check_file_type(path):
     elif os.path.isfile(path):
         return FileCheckConst.FILE
     else:
-        logger.error('Neither a file nor a directory.')
+        logger.error(f'{path} does not exist, please check!')
         raise FileCheckException(FileCheckException.INVALID_FILE_ERROR)
 
 
@@ -338,10 +368,10 @@ def load_yaml(yaml_path):
     return yaml_data
 
 
-def load_npy(filepath, enable_pickle=False):
+def load_npy(filepath):
     check_file_or_directory_path(filepath)
     try:
-        npy = np.load(filepath, allow_pickle=enable_pickle)
+        npy = np.load(filepath, allow_pickle=False)
     except Exception as e:
         logger.error(f"The numpy file failed to load. Please check the path: {filepath}.")
         raise RuntimeError(f"Load numpy file {filepath} failed.") from e
@@ -351,7 +381,7 @@ def load_npy(filepath, enable_pickle=False):
 def load_json(json_path):
     try:
         with FileOpen(json_path, "r") as f:
-            fcntl.flock(f, fcntl.LOCK_EX)
+            fcntl.flock(f, fcntl.LOCK_SH)
             data = json.load(f)
             fcntl.flock(f, fcntl.LOCK_UN)
     except Exception as e:
@@ -360,11 +390,11 @@ def load_json(json_path):
     return data
 
 
-def save_json(json_path, data, indent=None):
-    json_path = os.path.realpath(json_path)
+def save_json(json_path, data, indent=None, mode="w"):
     check_path_before_create(json_path)
+    json_path = os.path.realpath(json_path)
     try:
-        with FileOpen(json_path, 'w') as f:
+        with FileOpen(json_path, mode) as f:
             fcntl.flock(f, fcntl.LOCK_EX)
             json.dump(data, f, indent=indent)
             fcntl.flock(f, fcntl.LOCK_UN)
@@ -374,6 +404,35 @@ def save_json(json_path, data, indent=None):
     change_mode(json_path, FileCheckConst.DATA_FILE_AUTHORITY)
 
 
+def save_yaml(yaml_path, data):
+    check_path_before_create(yaml_path)
+    yaml_path = os.path.realpath(yaml_path)
+    try:
+        with FileOpen(yaml_path, 'w') as f:
+            fcntl.flock(f, fcntl.LOCK_EX)
+            yaml.dump(data, f, sort_keys=False)
+            fcntl.flock(f, fcntl.LOCK_UN)
+    except Exception as e:
+        logger.error(f'Save yaml file "{os.path.basename(yaml_path)}" failed.')
+        raise RuntimeError(f"Save yaml file {yaml_path} failed.") from e
+    change_mode(yaml_path, FileCheckConst.DATA_FILE_AUTHORITY)
+
+
+def save_excel(path, data):
+    check_path_before_create(path)
+    path = os.path.realpath(path)
+    try:
+        if isinstance(data, pd.DataFrame):
+            data.to_excel(path, index=False)
+        else:
+            logger.error(f'unsupported data type.')
+            return
+    except Exception as e:
+        logger.error(f'Save excel file "{os.path.basename(path)}" failed.')
+        raise RuntimeError(f"Save excel file {path} failed.") from e
+    change_mode(path, FileCheckConst.DATA_FILE_AUTHORITY)
+
+
 def move_file(src_path, dst_path):
     check_file_or_directory_path(src_path)
     check_path_before_create(dst_path)
@@ -386,8 +445,8 @@ def move_file(src_path, dst_path):
 
 
 def save_npy(data, filepath):
-    filepath = os.path.realpath(filepath)
     check_path_before_create(filepath)
+    filepath = os.path.realpath(filepath)
     try:
         np.save(filepath, data)
     except Exception as e:
@@ -396,9 +455,9 @@ def save_npy(data, filepath):
     change_mode(filepath, FileCheckConst.DATA_FILE_AUTHORITY)
 
 
-def save_npy_to_txt(self, data, dst_file='', align=0):
+def save_npy_to_txt(data, dst_file='', align=0):
     if os.path.exists(dst_file):
-        self.log.info("Dst file %s exists, will not save new one.", dst_file)
+        logger.info("Dst file %s exists, will not save new one." % dst_file)
         return
     shape = data.shape
     data = data.flatten()
@@ -408,10 +467,11 @@ def save_npy_to_txt(self, data, dst_file='', align=0):
         pad_array = np.zeros((align - data.size % align,))
         data = np.append(data, pad_array)
     check_path_before_create(dst_file)
+    dst_file = os.path.realpath(dst_file)
     try:
         np.savetxt(dst_file, data.reshape((-1, align)), delimiter=' ', fmt='%g')
     except Exception as e:
-        self.log.error("An unexpected error occurred: %s when savetxt to %s" % (str(e)), dst_file)
+        logger.error("An unexpected error occurred: %s when savetxt to %s" % (str(e), dst_file))
     change_mode(dst_file, FileCheckConst.DATA_FILE_AUTHORITY)
 
 
@@ -421,8 +481,8 @@ def save_workbook(workbook, file_path):
     workbook: 要保存的工作簿对象
     file_path: 文件保存路径
     """
-    file_path = os.path.realpath(file_path)
     check_path_before_create(file_path)
+    file_path = os.path.realpath(file_path)
     try:
         workbook.save(file_path)
     except Exception as e:
@@ -431,9 +491,27 @@ def save_workbook(workbook, file_path):
     change_mode(file_path, FileCheckConst.DATA_FILE_AUTHORITY)
 
 
-def write_csv(data, filepath, mode="a+"):
-    file_path = os.path.realpath(filepath)
+def write_csv(data, filepath, mode="a+", malicious_check=False):
+    def csv_value_is_valid(value: str) -> bool:
+        if not isinstance(value, str):
+            return True
+        try:
+            # -1.00 or +1.00 should be consdiered as digit numbers
+            float(value)
+        except ValueError:
+            # otherwise, they will be considered as formular injections
+            return not bool(re.compile(FileCheckConst.CSV_BLACK_LIST).search(value))
+        return True
+
+    if malicious_check:
+        for row in data:
+            for cell in row:
+                if not csv_value_is_valid(cell):
+                    raise RuntimeError(f"Malicious value [{cell}] is not allowed "
+                                       f"to be written into the csv: {filepath}.")
+
     check_path_before_create(filepath)
+    file_path = os.path.realpath(filepath)
     try:
         with FileOpen(filepath, mode, encoding='utf-8-sig') as f:
             writer = csv.writer(f)
@@ -444,6 +522,54 @@ def write_csv(data, filepath, mode="a+"):
     change_mode(filepath, FileCheckConst.DATA_FILE_AUTHORITY)
 
 
+def read_csv(filepath, as_pd=True):
+    check_file_or_directory_path(filepath)
+    try:
+        if as_pd:
+            csv_data = pd.read_csv(filepath)
+        else:
+            with FileOpen(filepath, 'r', encoding='utf-8-sig') as f:
+                csv_reader = csv.reader(f, delimiter=',')
+                csv_data = list(csv_reader)
+    except Exception as e:
+        logger.error(f"The csv file failed to load. Please check the path: {filepath}.")
+        raise RuntimeError(f"Read csv file {filepath} failed.") from e
+    return csv_data
+
+
+def write_df_to_csv(data, filepath, mode="w", header=True, malicious_check=False):
+    def csv_value_is_valid(value: str) -> bool:
+        if not isinstance(value, str):
+            return True
+        try:
+            # -1.00 or +1.00 should be consdiered as digit numbers
+            float(value)
+        except ValueError:
+            # otherwise, they will be considered as formular injections
+            return not bool(re.compile(FileCheckConst.CSV_BLACK_LIST).search(value))
+        return True
+
+    if not isinstance(data, pd.DataFrame):
+        raise ValueError("The data type of data is not supported. Only support pd.DataFrame.")
+
+    if malicious_check:
+        for i in range(len(data)):
+            for j in range(len(data.columns)):
+                cell = data.iloc[i, j]
+                if not csv_value_is_valid(cell):
+                    raise RuntimeError(f"Malicious value [{cell}] is not allowed "
+                                       f"to be written into the csv: {filepath}.")
+
+    check_path_before_create(filepath)
+    file_path = os.path.realpath(filepath)
+    try:
+        data.to_csv(filepath, mode=mode, header=header, index=False)
+    except Exception as e:
+        logger.error(f'Save csv file "{os.path.basename(file_path)}" failed')
+        raise RuntimeError(f"Save csv file {file_path} failed.") from e
+    change_mode(filepath, FileCheckConst.DATA_FILE_AUTHORITY)
+
+
 def remove_path(path):
     if not os.path.exists(path):
         return
@@ -476,3 +602,46 @@ def get_json_contents(file_path):
 def get_file_content_bytes(file):
     with FileOpen(file, 'rb') as file_handle:
         return file_handle.read()
+
+
+# 对os.walk设置遍历深度
+def os_walk_for_files(path, depth):
+    res = []
+    for root, _, files in os.walk(path, topdown=True):
+        check_file_or_directory_path(root, isdir=True)
+        if root.count(os.sep) - path.count(os.sep) >= depth:
+            _[:] = []
+        else:
+            for file in files:
+                res.append({"file": file, "root": root})
+    return res
+
+
+def check_crt_valid(pem_path):
+    """
+    Check the validity of the SSL certificate.
+
+    Load the SSL certificate from the specified path, parse and check its validity period.
+    If the certificate is expired or invalid, raise a RuntimeError.
+
+    Parameters:
+    pem_path (str): The file path of the SSL certificate.
+
+    Raises:
+    RuntimeError: If the SSL certificate is invalid or expired.
+    """
+    try:
+        with FileOpen(pem_path, "r") as f:
+            pem_data = f.read()
+        cert = OpenSSL.crypto.load_certificate(OpenSSL.crypto.FILETYPE_PEM, pem_data)
+        pem_start = parser.parse(cert.get_notBefore().decode("UTF-8"))
+        pem_end = parser.parse(cert.get_notAfter().decode("UTF-8"))
+        logger.info(f"The SSL certificate passes the verification and the validity period "
+                    f"starts from {pem_start} ends at {pem_end}.")
+    except Exception as e:
+        logger.error("Failed to parse the SSL certificate. Check the certificate.")
+        raise RuntimeError(f"The SSL certificate is invalid, {pem_path}") from e
+
+    now_utc = datetime.now(tz=timezone.utc)
+    if cert.has_expired() or not (pem_start <= now_utc <= pem_end):
+        raise RuntimeError(f"The SSL certificate has expired and needs to be replaced, {pem_path}")

msprobe/core/common/inplace_op_checker.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from msprobe.core.common.file_utils import load_yaml
+
+
+class InplaceOpChecker:
+    OP_FUNCTIONAL = 'functional'
+    OP_TENSOR = 'tensor'
+    OP_TORCH = 'torch'
+    OP_DISTRIBUTED = 'distributed'
+
+    INPLACE_OPS_DICT = None
+
+    @classmethod
+    def load_ops(cls):
+        if cls.INPLACE_OPS_DICT is None:
+            cls.INPLACE_OPS_DICT = dict()
+        cur_path = os.path.dirname(os.path.realpath(__file__))
+        yaml_path = os.path.join(cur_path, "inplace_ops.yaml")
+        all_ops = load_yaml(yaml_path)
+        cls.INPLACE_OPS_DICT[cls.OP_FUNCTIONAL] = all_ops.get('inplace_functional_op')
+        cls.INPLACE_OPS_DICT[cls.OP_TENSOR] = all_ops.get('inplace_tensor_op')
+        cls.INPLACE_OPS_DICT[cls.OP_TORCH] = all_ops.get('inplace_torch_op')
+        cls.INPLACE_OPS_DICT[cls.OP_DISTRIBUTED] = all_ops.get('inplace_distributed_op')
+
+    @classmethod
+    def check(cls, api, category='distributed'):
+        """
+            给定api和分类，检查其是否为inplace操作
+        """
+        if not cls.INPLACE_OPS_DICT:
+            cls.load_ops()
+
+        if category not in cls.INPLACE_OPS_DICT.keys():
+            return False
+        return api in cls.INPLACE_OPS_DICT[category]
+
+
+InplaceOpChecker.load_ops()

msprobe/core/common/inplace_ops.yaml

@@ -0,0 +1,251 @@
+inplace_functional_op:
+  - threshold_
+  - relu_
+  - hardtanh_
+  - elu_
+  - selu_
+  - celu_
+  - leaky_relu_
+  - rrelu_
+
+inplace_tensor_op:
+  - __iadd__
+  - __iand__
+  - __idiv__
+  - __ifloordiv__
+  - __ilshift__
+  - __imod__
+  - __imul__
+  - __ior__
+  - __irshift__
+  - __isub__
+  - __ixor__
+  - abs_
+  - absolute_
+  - acos_
+  - acosh_
+  - add_
+  - addbmm_
+  - addcdiv_
+  - addcmul_
+  - addmm_
+  - addmv_
+  - addr_
+  - arccos_
+  - arccosh_
+  - arcsin_
+  - arcsinh_
+  - arctan_
+  - arctanh_
+  - asin_
+  - asinh_
+  - atan2_
+  - atan_
+  - atanh_
+  - baddbmm_
+  - bernoulli_
+  - bitwise_and_
+  - bitwise_not_
+  - bitwise_or_
+  - bitwise_xor_
+  - cauchy_
+  - ceil_
+  - clamp_
+  - clamp_max_
+  - clamp_min_
+  - clip_
+  - copysign_
+  - cos_
+  - cosh_
+  - cumprod_
+  - cumsum_
+  - deg2rad_
+  - digamma_
+  - div_
+  - divide_
+  - eq_
+  - erf_
+  - erfc_
+  - erfinv_
+  - exp2_
+  - exp_
+  - expm1_
+  - exponential_
+  - fill_
+  - fill_diagonal_
+  - fix_
+  - float_power_
+  - floor_
+  - floor_divide_
+  - fmod_
+  - frac_
+  - gcd_
+  - ge_
+  - geometric_
+  - greater_
+  - gt_
+  - greater_equal_
+  - heaviside_
+  - hypot_
+  - igamma_
+  - igammac_
+  - index_add_
+  - index_copy_
+  - index_fill_
+  - index_put_
+  - lcm_
+  - ldexp_
+  - le_
+  - lerp_
+  - less_
+  - less_equal_
+  - lgamma_
+  - log10_
+  - log1p_
+  - log2_
+  - log_
+  - log_normal_
+  - logical_and_
+  - logical_not_
+  - logical_or_
+  - logical_xor_
+  - logit_
+  - lt_
+  - map2_
+  - map_
+  - masked_fill_
+  - masked_scatter_
+  - mul_
+  - multiply_
+  - mvlgamma_
+  - ne_
+  - neg_
+  - negative_
+  - normal_
+  - not_equal_
+  - pow_
+  - polygamma_
+  - put_
+  - rad2deg_
+  - reciprocal_
+  - relu_
+  - remainder_
+  - renorm_
+  - resize_
+  - resize_as_
+  - round_
+  - rsqrt_
+  - scatter_
+  - scatter_add_
+  - sgn_
+  - sigmoid_
+  - sign_
+  - sin_
+  - sinc_
+  - sinh_
+  - sqrt_
+  - square_
+  - squeeze_
+  - sub_
+  - t_
+  - tan_
+  - tanh_
+  - transpose_
+  - tril_
+  - triu_
+  - true_divide_
+  - trunc_
+  - unsqueeze_
+  - xlogy_
+
+inplace_torch_op:
+  - _add_relu_
+  - abs_
+  - acos_
+  - acosh_
+  - addmv_
+  - alpha_dropout_
+  - arccos_
+  - arccosh_
+  - arcsin_
+  - arcsinh_
+  - arctan_
+  - arctanh_
+  - asin_
+  - asinh_
+  - atan_
+  - atanh_
+  - ceil_
+  - celu_
+  - clamp_
+  - clamp_max_
+  - clamp_min_
+  - clip_
+  - cos_
+  - cosh_
+  - deg2rad_
+  - dropout_
+  - embedding_renorm_
+  - erf_
+  - erfc_
+  - exp2_
+  - exp_
+  - expm1_
+  - feature_alpha_dropout_
+  - feature_dropout_
+  - fill_
+  - fix_
+  - floor_
+  - frac_
+  - gcd_
+  - index_put_
+  - lcm_
+  - ldexp_
+  - log10_
+  - log1p_
+  - log2_
+  - log_
+  - logit_
+  - nan_to_num_
+  - neg_
+  - negative_
+  - rad2deg_
+  - reciprocal_
+  - relu_
+  - resize_as_
+  - round_
+  - rrelu_
+  - rsqrt_
+  - selu_
+  - sigmoid_
+  - sin_
+  - sinc_
+  - sinh_
+  - sqrt_
+  - square_
+  - tan_
+  - tanh_
+  - threshold_
+  - trunc_
+  - xlogy_
+
+inplace_distributed_op:
+  - broadcast
+  - all_reduce
+  - reduce
+  - all_gather
+  - gather
+  - scatter
+  - reduce_scatter
+  - _reduce_scatter_base
+  - _all_gather_base
+  - send
+  - recv
+  - irecv
+  - isend
+  - all_to_all_single
+  - all_to_all
+  - all_gather_into_tensor
+  - reduce_scatter_tensor
+
+