maya-umbrella 0.4.1__py2.py3-none-any.whl → 0.6.0__py2.py3-none-any.whl

maya_umbrella/defender.py

@@ -0,0 +1,161 @@
+# Import built-in modules
+from contextlib import contextmanager
+import logging
+
+# Import local modules
+from maya_umbrella.cleaner import MayaVirusCleaner
+from maya_umbrella.collector import MayaVirusCollector
+from maya_umbrella.filesystem import get_hooks
+from maya_umbrella.filesystem import load_hook
+from maya_umbrella.i18n import Translator
+from maya_umbrella.log import setup_logger
+from maya_umbrella.maya_funs import is_maya_standalone
+from maya_umbrella.maya_funs import om
+
+
+# Global list to store IDs of Maya callbacks
+MAYA_UMBRELLA_CALLBACK_IDS = []
+
+
+def _add_callbacks_id(id_):
+    """Add a callback ID to the global list if it's not already present.
+
+    Args:
+        id_ (int): ID of the callback to be added.
+    """
+    global MAYA_UMBRELLA_CALLBACK_IDS
+    if id_ not in MAYA_UMBRELLA_CALLBACK_IDS:
+        MAYA_UMBRELLA_CALLBACK_IDS.append(id_)
+
+
+class MayaVirusDefender(object):
+    """A class to defend against Maya viruses.
+
+    Attributes:
+        _vaccines (list): List to store vaccines.
+        callback_maps (dict): Dictionary to map callback names to MSceneMessage constants.
+        auto_fix (bool): Whether to automatically fix issues.
+        logger (Logger): Logger object for logging purposes.
+        translator (Translator): Translator object for translation purposes.
+        collector (MayaVirusCollector): MayaVirusCollector object for collecting issues.
+        virus_cleaner (MayaVirusCleaner): MayaVirusCleaner object for fixing issues.
+        hooks (list): List of hooks to run.
+    """
+    _vaccines = []
+    callback_maps = {
+        "after_open": om.MSceneMessage.kAfterOpen,
+        "maya_initialized": om.MSceneMessage.kMayaInitialized,
+        "after_import": om.MSceneMessage.kAfterImport,
+        "after_import_reference": om.MSceneMessage.kAfterImportReference,
+        "after_load_reference": om.MSceneMessage.kAfterLoadReference,
+        "before_save": om.MSceneMessage.kBeforeSave,
+        "before_import": om.MSceneMessage.kBeforeImport,
+        "before_load_reference": om.MSceneMessage.kBeforeLoadReference,
+        "before_import_reference": om.MSceneMessage.kBeforeImportReference,
+        "maya_exiting": om.MSceneMessage.kMayaExiting,
+    }
+
+    def __init__(self, auto_fix=True):
+        """Initialize the MayaVirusDefender.
+
+        Args:
+            auto_fix (bool): Whether to automatically fix issues.
+        """
+        logger = logging.getLogger(__name__)
+        self.auto_fix = auto_fix
+        self.logger = setup_logger(logger)
+        self.translator = Translator()
+        self.collector = MayaVirusCollector(self.logger, self.translator)
+        self.virus_cleaner = MayaVirusCleaner(self.collector, self.logger)
+        self.hooks = get_hooks()
+
+    def run_hooks(self):
+        """Run all hooks, only works in non-batch mode."""
+        if not is_maya_standalone():
+            for hook_file in self.hooks:
+                self.logger.debug("run_hook: %s", hook_file)
+                try:
+                    load_hook(hook_file).hook(virus_cleaner=self.virus_cleaner)
+                except Exception as e:
+                    self.logger.debug("Error running hook: %s", e)
+
+    def collect(self):
+        """Collect all issues related to the Maya virus."""
+        self.collector.collect()
+
+    def fix(self):
+        """Fix all issues related to the Maya virus."""
+        self.virus_cleaner.fix()
+
+    def report(self):
+        """Report all issues related to the Maya virus."""
+        self.collect()
+        self.collector.report()
+
+    @property
+    def have_issues(self):
+        """Check if any issues are found.
+
+        Returns:
+            bool: True if any issues are found, False otherwise.
+        """
+        return self.collector.have_issues
+
+    def setup(self):
+        """Set up the MayaVirusDefender."""
+        self.virus_cleaner.setup_default_callbacks()
+        for name, callbacks in self.collector.registered_callbacks.items():
+            maya_callback = self.callback_maps[name]
+            self.logger.debug("%s setup.", name)
+            for func in callbacks:
+                _add_callbacks_id(om.MSceneMessage.addCallback(maya_callback, func))
+        for name, callbacks in self.callback_maps.items():
+            self.logger.debug("setup callback %s.", name)
+            _add_callbacks_id(om.MSceneMessage.addCallback(callbacks, self._callback))
+
+    def stop(self):
+        """Stop the MayaVirusDefender."""
+        for ids in MAYA_UMBRELLA_CALLBACK_IDS:
+            self.logger.debug("remove callback. %s", ids)
+            om.MSceneMessage.removeCallback(ids)
+            MAYA_UMBRELLA_CALLBACK_IDS.remove(ids)
+
+    def get_unfixed_references(self):
+        """Get the list of unfixed reference files.
+
+        Returns:
+            list: List of unfixed reference files.
+        """
+        self.collect()
+        return self.collector.infected_reference_files
+
+    def _callback(self, *args, **kwargs):
+        """Callback function for MayaVirusDefender.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        if self.auto_fix:
+            self.collect()
+            self.fix()
+            self.run_hooks()
+        else:
+            self.report()
+
+    def start(self):
+        """Start the MayaVirusDefender."""
+        self._callback()
+
+
+@contextmanager
+def context_defender():
+    """Context manager for MayaVirusDefender.
+
+    Yields:
+        MayaVirusDefender: An instance of MayaVirusDefender.
+    """
+    defender = MayaVirusDefender()
+    defender.stop()
+    yield defender
+    defender.setup()

maya_umbrella/filesystem.py

@@ -1,17 +1,25 @@
 # Import built-in modules
+from contextlib import contextmanager
 import glob
 import importlib
-import io
+import json
 import os
 import random
+import re
 import shutil
 import string
+import sys
 import tempfile
 
 # Import local modules
+from maya_umbrella.constants import FILE_VIRUS_SIGNATURES
 from maya_umbrella.constants import PACKAGE_NAME
 
 
+PY2 = sys.version_info[0] == 2
+PY3 = sys.version_info[0] == 3
+
+
 def this_root():
     """Return the absolute path of the current file's directory."""
     return os.path.abspath(os.path.dirname(__file__))
@@ -21,7 +29,7 @@ def safe_remove_file(file_path):
     """Remove the file at the given path without raising an error if the file does not exist."""
     try:
         os.remove(file_path)
-    except OSError:
+    except (OSError, IOError):  # noqa: UP024
         pass
 
 
@@ -29,42 +37,65 @@ def safe_rmtree(path):
     """Remove the directory at the given path without raising an error if the directory does not exist."""
     try:
         shutil.rmtree(path)
-    except OSError:
+    except (OSError, IOError):  # noqa: UP024
         pass
 
 
 def read_file(path):
     """Read the content of the file at the given path."""
-    with io.open(path, "r", encoding="utf-8") as file_:
-        content = file_.read()
+    options = {"encoding": "utf-8"} if PY3 else {}
+    with open(path, **options) as file_:
+        try:
+            content = file_.read()
+        # maya-2022 UnicodeDecodeError from `plug-ins/mayaHIK.pres.mel`
+        except UnicodeDecodeError:
+            return ""
+    return content
+
+
+def read_json(path):
+    """Read the content of the file at the given path."""
+    options = {"encoding": "utf-8"} if PY3 else {}
+    with open(path, **options) as file_:
+        try:
+            content = json.load(file_)
+        except UnicodeDecodeError:
+            return {}
     return content
 
 
 def write_file(path, content):
     """Write the given content to the file at the given path."""
-    with io.open(path, "w", encoding="utf-8", newline="\n") as file_:
+    options = {"encoding": "utf-8"} if PY3 else {}
+    with atomic_writes(path, "w", **options) as file_:
         file_.write(content)
 
 
-def patch_file(source, target, key_values, report_error=True):
-    """Modify the file at the source path with the given key value pairs and write the result to the target path.
+@contextmanager
+def atomic_writes(src, mode, **options):
+    """Context manager for atomic writes to a file.
 
-    If report_error is True, raise an IndexError if any of the keys are not found in the source file.
+    This context manager ensures that the file is only written to disk if the write operation completes without errors.
+
+    Args:
+        src (str): Path to the file to be written.
+        mode (str): Mode in which the file is opened, like 'r', 'w', 'a', etc.
+        **options: Arbitrary keyword arguments that are passed to the built-in open() function.
+
+    Yields:
+        file object: The opened file object.
+
+    Raises:
+        AttributeError: If the os module does not have the 'replace' function (Python 2 compatibility).
     """
-    key_values = key_values if key_values else {}
-    found_keys = []
-    file_data = read_file(source)
-    for key, value in key_values.items():
-        before_ = file_data
-        file_data = file_data.replace(key, value)
-        if before_ != file_data:
-            found_keys.append(key)
-    write_file(target, file_data)
-
-    if report_error:
-        not_found = list(set(key_values.keys()) - set(found_keys))
-        if not_found:
-            raise IndexError("Not found: {0}".format(not_found))
+    temp_path = os.path.join(os.path.dirname(src), "._{}".format(id_generator()))
+    with open(temp_path, mode, **options) as f:
+        yield f
+    try:
+        os.replace(temp_path, src)
+    except AttributeError:
+        shutil.move(temp_path, src)
+
 
 
 def id_generator(size=6, chars=string.ascii_uppercase + string.digits):
@@ -77,7 +108,7 @@ def rename(src):
     dst = os.path.join(os.path.dirname(src), "._{}".format(id_generator()))
     try:
         os.rename(src, dst)
-    except OSError:
+    except (OSError, IOError):  # noqa: UP024
         return src
     return dst
 
@@ -88,7 +119,7 @@ def load_hook(hook_file):
     if hasattr(importlib, "machinery"):
         # Python 3
         # Import built-in modules
-        from importlib.util import spec_from_loader
+        from importlib.util import spec_from_loader  # noqa: F401
 
         loader = importlib.machinery.SourceFileLoader(hook_name, hook_file)
         spec = importlib.util.spec_from_loader(loader.name, loader=loader)
@@ -137,7 +168,101 @@ def get_log_file():
     root = get_log_root()
     try:
         os.makedirs(root)
-    except OSError:
+    except (OSError, IOError):  # noqa: UP024
         pass
     name = os.getenv("MAYA_UMBRELLA_LOG_NAME", PACKAGE_NAME)
     return os.path.join(root, "{name}.log".format(name=name))
+
+
+def remove_virus_file_by_signature(file_path, signatures, output_file_path=None):
+    """Remove virus content from a file by matching signatures.
+
+    Args:
+        file_path (str): Path to the file to be cleaned.
+        signatures (list): List of signatures to match and remove.
+        output_file_path (str, optional): Path to the cleaned output file.
+         Defaults to None, which overwrites the input file.
+    """
+    data = read_file(file_path)
+    if check_virus_by_signature(data, signatures):
+        fixed_data = replace_content_by_signatures(data, signatures)
+        write_file(output_file_path or file_path, fixed_data)
+
+
+def replace_content_by_signatures(content, signatures):
+    """Replace content in a string that matches given signatures.
+
+    Args:
+        content (str): The input content.
+        signatures (list): List of signatures to match and remove.
+
+    Returns:
+        str: The cleaned content.
+    """
+    for signature in signatures:
+        content = re.sub(signature, "", content)
+    return content
+
+
+def check_virus_file_by_signature(file_path, signatures=None):
+    """Check if a file contains a virus by matching signatures.
+
+    Args:
+        file_path (str): Path to the file to be checked.
+        signatures (list, optional): List of signatures to match. Defaults to None, which uses FILE_VIRUS_SIGNATURES.
+
+    Returns:
+        bool: True if a virus signature is found, False otherwise.
+    """
+    signatures = signatures or FILE_VIRUS_SIGNATURES
+    try:
+        data = read_file(file_path)
+        return check_virus_by_signature(data, signatures)
+    except (OSError, IOError):  # noqa: UP024
+        return False
+    except UnicodeDecodeError:
+        return True
+
+
+def check_virus_by_signature(content, signatures=None):
+    """Check if a content contains a virus by matching signatures.
+
+    Args:
+        content (str): The input content.
+        signatures (list, optional): List of signatures to match. Defaults to None, which uses FILE_VIRUS_SIGNATURES.
+
+    Returns:
+        bool: True if a virus signature is found, False otherwise.
+    """
+    signatures = signatures or FILE_VIRUS_SIGNATURES
+    for signature in signatures:
+        if re.search(signature, content):
+            return True
+    return False
+
+
+def get_backup_path(path, root_path=None):
+    """Get the backup path for a given file path based on environment variables.
+
+    Args:
+        path (str): Path to the original file.
+        root_path (str, optional): Path to the root folder where backups should be saved.
+            Defaults to None, which saves backups in the original file's folder.
+
+    Returns:
+        str: The backup path.
+    """
+    ignore_backup = os.getenv("MAYA_UMBRELLA_IGNORE_BACKUP", "false").lower() == "true"
+    if ignore_backup:
+        return path
+    root, filename = os.path.split(path)
+    backup_folder_name = os.getenv("MAYA_UMBRELLA_BACKUP_FOLDER_NAME", "_virus")
+    backup_path = os.path.join(root, backup_folder_name)
+    if root_path:
+        _, base_path = os.path.splitdrive(root)
+        backup_path = os.path.join(root_path, base_path.strip(os.sep))
+    try:
+        os.makedirs(backup_path)
+    except (OSError, IOError):  # noqa: UP024
+        pass
+    return os.path.join(backup_path, filename)

maya_umbrella/hooks/delete_turtle.py

@@ -1,6 +1,6 @@
-# Import third-party modules
-import maya.cmds as cmds
-import maya.mel as mel
+# Import local modules
+from maya_umbrella.maya_funs import cmds
+from maya_umbrella.maya_funs import mel
 
 
 def hook(virus_cleaner):

maya_umbrella/hooks/delete_unknown_plugin_node.py

@@ -1,5 +1,5 @@
-# Import third-party modules
-import maya.cmds as cmds
+# Import local modules
+from maya_umbrella.maya_funs import cmds
 
 
 def hook(virus_cleaner):
@@ -16,7 +16,8 @@ def hook(virus_cleaner):
                         cmds.lockNode(nodeObj, lock=False)
                     except Exception:
                         virus_cleaner.logger.warning(
-                            "The node is locked and cannot be unlocked. skip {}".format(nodeObj))
+                            "The node is locked and cannot be unlocked. skip {}".format(nodeObj)
+                        )
                         continue
                 try:
                     cmds.delete(nodeObj)

maya_umbrella/hooks/fix_model_panel.py

@@ -1,5 +1,5 @@
-# Import third-party modules
-import maya.cmds as cmds
+# Import local modules
+from maya_umbrella.maya_funs import cmds
 
 
 def hook(virus_cleaner):

maya_umbrella/hooks/fix_on_model_change_3dc.py

@@ -1,6 +1,6 @@
-# Import third-party modules
-import maya.cmds as cmds
-import maya.mel as mel
+# Import local modules
+from maya_umbrella.maya_funs import cmds
+from maya_umbrella.maya_funs import mel
 
 
 def hook(virus_cleaner):

maya_umbrella/i18n.py

@@ -0,0 +1,76 @@
+# Import built-in modules
+import glob
+import os
+from string import Template
+
+# Import local modules
+from maya_umbrella.filesystem import read_json
+from maya_umbrella.filesystem import this_root
+from maya_umbrella.maya_funs import maya_ui_language
+
+
+class Translator(object):
+    """A class to handle translations for different locales.
+
+    Attributes:
+        data (dict): Dictionary containing translation data for different locales.
+        locale (str): The current locale.
+    """
+    def __init__(self, file_format="json", default_locale=None):
+        """Initialize the Translator.
+
+        Args:
+            file_format (str, optional): File format of the translation files. Defaults to "json".
+            default_locale (str, optional): Default locale to use for translations. Defaults to None,
+            which uses the MAYA_UMBRELLA_LANG environment variable or the Maya UI language.
+        """
+        _default_locale = os.getenv("MAYA_UMBRELLA_LANG", maya_ui_language())
+        default_locale = default_locale or _default_locale
+        self.data = {}
+        self.locale = default_locale
+        translations_folder = os.path.join(this_root(), "locales")
+
+        # get list of files with specific extensions
+        files = glob.glob(os.path.join(translations_folder, "*.{file_format}".format(file_format=file_format)))
+        for fil in files:
+            # get the name of the file without extension, will be used as locale name
+            loc = os.path.splitext(os.path.basename(fil))[0]
+            self.data[loc] = read_json(fil)
+
+    def set_locale(self, locale):
+        """Set the current locale.
+
+        Args:
+            locale (str): The locale to set.
+
+        Raises:
+            ValueError: If the provided locale is not supported.
+        """
+        if locale in self.data:
+            self.locale = locale
+        else:
+            raise ValueError("Invalid locale: {loc}".format(loc=locale))
+
+    def get_locale(self):
+        """Get the current locale.
+
+        Returns:
+            str: The current locale.
+        """
+        return self.locale
+
+    def translate(self, key, **kwargs):
+        """Translate a text based on the current locale.
+
+        Args:
+            key (str): The key to be translated.
+            **kwargs: Arbitrary keyword arguments that are used to replace placeholders in the translation text.
+
+        Returns:
+            str: The translated text.
+        """
+        # return the key instead of translation text if locale is not supported
+        if self.locale not in self.data:
+            return key
+        text = self.data[self.locale].get(key, key)
+        return Template(text).safe_substitute(**kwargs)

maya_umbrella/locales/en_US.json

@@ -0,0 +1,19 @@
+{
+  "start_fix_issues": "Start fixing all problems related to Maya virus $name",
+  "finish_fix_issues": "Done.",
+  "init_message": "Successfully loaded <hl>maya_umbrella</hl> under protection.",
+  "init_standalone_message": "-----------------------Loading maya_umbrella successfully----------------------",
+  "report_issue": "$name: Infected by Malware!",
+  "infected_nodes": "Infected nodes: $name: ",
+  "bad_files": "Bad files: $name",
+  "infected_script_jobs": "Infected script jobs: $name",
+  "infected_files": "Infected files: $name",
+  "infected_reference_files": "Infected reference files: $name",
+  "fix_infected_files": "Clean infected files: $name",
+  "fix_infected_nodes": "Delete infected nodes：$name",
+  "fix_infected_reference_nodes": "trying fix infected nodes from reference：$name",
+  "delete": "Deleting: $name",
+  "remove_file": "Deleting file：$name",
+  "remove_path": "Deleting path：$name",
+  "fix_script_job": "Kill script job: %s"
+}

maya_umbrella/locales/zh_CN.json

@@ -0,0 +1,19 @@
+{
+  "start_fix_issues": "开始修复与 Maya 病毒相关的所有问题 $name",
+  "finish_fix_issues": "修复结束.",
+  "init_message": "成功加载 <hl>maya_umbrella</hl> 保护中.",
+  "init_standalone_message": "-----------------------成功加载 maya_umbrella-----------------------",
+  "report_issue": "$name：被恶意感染！",
+  "infected_nodes": "被感染节点：$name: ",
+  "bad_files": "需要被清理的文件：$name",
+  "infected_script_jobs": "被感染的script jobs：$name",
+  "infected_files": "被感染的脚本：$name",
+  "infected_reference_files": "被感染的参考文件：$name",
+  "fix_infected_files": "清理被感染的文件：$name",
+  "fix_infected_nodes": "删除被感染的节点：$name",
+  "fix_infected_reference_nodes": "尝试修复被感染的参考节点：$name",
+  "delete": "删除感染文件：$name",
+  "remove_file": "删除文件：$name",
+  "remove_path": "删除文件夹：$name",
+  "fix_script_job": "删除被感染的节点：$name"
+}