maya-umbrella 0.5.0__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

@@ -2,7 +2,7 @@
 from contextlib import contextmanager
 import glob
 import importlib
-import logging
+import json
 import os
 import random
 import re
@@ -53,6 +53,17 @@ def read_file(path):
     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."""
     options = {"encoding": "utf-8"} if PY3 else {}
@@ -62,6 +73,21 @@ def write_file(path, content):
 
 @contextmanager
 def atomic_writes(src, mode, **options):
+    """Context manager for atomic writes to a 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).
+    """
     temp_path = os.path.join(os.path.dirname(src), "._{}".format(id_generator()))
     with open(temp_path, mode, **options) as f:
         yield f
@@ -71,6 +97,7 @@ def atomic_writes(src, mode, **options):
         shutil.move(temp_path, src)
 
 
+
 def id_generator(size=6, chars=string.ascii_uppercase + string.digits):
     """Generate a random string of the given size using the given characters."""
     return "".join(random.choice(chars) for _ in range(size))
@@ -148,21 +175,45 @@ def get_log_file():
 
 
 def remove_virus_file_by_signature(file_path, signatures, output_file_path=None):
-    logger = logging.getLogger(__name__)
+    """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):
-        logger.warning("%s: Infected by Malware!", file_path)
         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)
@@ -174,8 +225,44 @@ def check_virus_file_by_signature(file_path, signatures=None):
 
 
 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/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"
+}

maya_umbrella/maya_funs.py

@@ -22,15 +22,115 @@ except ImportError:
     om = MagicMock()
     mel = MagicMock()
 
+# Import built-in modules
+from functools import wraps
+
 
 def is_maya_standalone():
-    """Return True if Maya is standalone."""
+    """Check if Maya is running in standalone mode.
+
+    Returns:
+        bool: True if Maya is running in standalone mode, False otherwise.
+    """
     return cmds.about(batch=True)
 
 
 def check_reference_node_exists(node_name):
-    """Check if reference node exists."""
+    """Check if a reference node exists in the Maya scene.
+
+    Args:
+        node_name (str): Name of the reference node.
+
+    Returns:
+        bool: True if the reference node exists, False otherwise.
+    """
     try:
         return cmds.referenceQuery(node_name, isNodeReferenced=True)
     except RuntimeError:
         return False
+
+
+def get_reference_file_by_node(node_name):
+    """Get the reference file associated with a node.
+
+    Args:
+        node_name (str): Name of the node.
+
+    Returns:
+        str: Path of the reference file, empty string if the node is not associated with a reference file.
+    """
+    try:
+        return cmds.referenceQuery(node_name, filename=True)
+    except RuntimeError:
+        return ""
+
+
+def get_attr_value(node_name, attr_name):
+    """Get the value of an attribute of a node.
+
+    Args:
+        node_name (str): Name of the node.
+        attr_name (str): Name of the attribute.
+
+    Returns:
+        Any: Value of the attribute, None if the attribute does not exist.
+    """
+    try:
+        return cmds.getAttr("{node_name}.{attr}".format(node_name=node_name, attr=attr_name))
+    except ValueError:
+        return None
+
+
+def maya_ui_language():
+    """Get the language of the Maya user interface.
+
+    Returns:
+        str: The language of the Maya user interface.
+    """
+    return cmds.about(uiLocaleLanguage=True)
+
+
+def block_prompt(func):
+    """Decorator to block file prompt dialogs in Maya.
+
+    Args:
+        func (function): The function to decorate.
+
+    Returns:
+        function: The decorated function.
+    """
+    @wraps(func)
+    def wrap(*args, **kwargs):
+        # Grabs the initial value.
+        prompt_val = cmds.file(prompt=True, q=True)
+
+        try:
+            cmds.file(prompt=False)
+            return func(*args, **kwargs)
+
+        finally:
+            # Resets to the original value, this way you don't suddenly turn the prompt on, when someone wanted it off.
+            cmds.file(prompt=prompt_val)
+
+    return wrap
+
+
+@block_prompt
+def open_maya_file(maya_file):
+    """Open a Maya file.
+
+    Args:
+        maya_file (str): Path to the Maya file.
+    """
+    cmds.file(maya_file, open=True, force=True, ignoreVersion=True, executeScriptNodes=False)
+
+
+@block_prompt
+def save_as_file(file_name):
+    """Save the current Maya scene as a file.
+
+    Args:
+        file_name (str): Path to the output file.
+    """
+    cmds.file(rename=file_name)
+    cmds.file(s=True, f=True)

maya_umbrella/scanner.py

@@ -0,0 +1,101 @@
+# Import built-in modules
+import glob
+import logging
+import os
+import shutil
+
+# Import local modules
+from maya_umbrella import maya_funs
+from maya_umbrella.defender import context_defender
+from maya_umbrella.filesystem import get_backup_path
+from maya_umbrella.filesystem import read_file
+from maya_umbrella.maya_funs import cmds
+
+
+class MayaVirusScanner(object):
+    """A class to scan and fix Maya files containing viruses.
+
+    Attributes:
+        _failed_files (list): List of files that failed to be fixed.
+        _fixed_files (list): List of files that have been fixed.
+        logger (Logger): Logger object for logging purposes.
+        defender (MayaVirusDefender): MayaVirusDefender object for fixing issues.
+        _env (dict): Custom environment variables.
+        output_path (str, optional): Path to save the fixed files. Defaults to None, which overwrites the original
+            files.
+    """
+
+    def __init__(self, output_path=None, env=None):
+        """Initialize the MayaVirusScanner.
+
+        Args:
+            output_path (str, optional): Path to save the fixed files. Defaults to None, which overwrites the original
+            files.
+            env (dict, optional): Custom environment variables. Defaults to None,
+            which sets the 'MAYA_COLOR_MANAGEMENT_SYNCOLOR' variable to '1'.
+        """
+        self.logger = logging.getLogger(__name__)
+        self.defender = None
+        self.output_path = output_path
+        self._failed_files = []
+        self._fixed_files = []
+        # Custom env.
+        self._env = env or {
+            "MAYA_COLOR_MANAGEMENT_SYNCOLOR": "1"
+        }
+
+    def scan_files_from_pattern(self, pattern):
+        """Scan and fix Maya files matching a given pattern.
+
+        Args:
+            pattern (str): The file pattern to match.
+        """
+        os.environ.update(self._env)
+        return self.scan_files_from_list(glob.iglob(pattern))
+
+    def scan_files_from_list(self, files):
+        """Scan and fix Maya files from a given list.
+
+        Args:
+            files (list): List of file paths to scan and fix.
+        """
+        with context_defender() as defender:
+            self.defender = defender
+            for maya_file in files:
+                self._fix(maya_file)
+        return self._fixed_files
+
+    def scan_files_from_file(self, text_file):
+        """Scan and fix Maya files from a given text file containing a list of file paths.
+
+        Args:
+            text_file (str): Path to the text file containing the list of file paths.
+        """
+        file_data = read_file(text_file)
+        files = file_data.splitlines()
+        return self.scan_files_from_list(files)
+
+    def _fix(self, maya_file):
+        """Fix a single Maya file containing a virus.
+
+        Args:
+            maya_file (str): Path to the Maya file to be fixed.
+        """
+        if not maya_file and maya_file in self._fixed_files:
+            self.logger.debug("Already fixed: {maya_file}".format(maya_file=maya_file))
+            return
+        try:
+            maya_funs.open_maya_file(maya_file)
+            self.defender.collect()
+        except Exception:
+            self._failed_files.append(maya_file)
+        if self.defender.have_issues:
+            self.defender.fix()
+            backup_path = get_backup_path(maya_file, root_path=self.output_path)
+            self.logger.debug("Backup saved to: {backup_path}".format(backup_path=backup_path))
+            shutil.copy2(maya_file, backup_path)
+            cmds.file(s=True, f=True)
+            self._fixed_files.append(maya_file)
+            cmds.file(new=True, force=True)
+        for ref in self.defender.collector.infected_reference_files:
+            self._fix(ref)

maya_umbrella/signatures.py

@@ -0,0 +1,10 @@
+# Import built-in modules
+from collections import namedtuple
+
+
+VirusSignature = namedtuple("VirusSignature", ["name", "signature"])
+
+# https://regex101.com/r/0MNzF7/1
+virus20240430_sig1 = VirusSignature("virus20240430", "python(.*);.+exec.+(pyCode).+;")
+# https://regex101.com/r/2D14UA/1
+virus20240430_sig2 = VirusSignature("virus20240430", r"^\['.+']")