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

maya_umbrella/maya_funs.py

@@ -0,0 +1,136 @@
+"""In order to facilitate unit testing in CI.
+
+we extracted all the interfaces used in Maya and mocked them uniformly.
+
+"""
+
+# Import third-party modules
+try:
+    # Import third-party modules
+    import maya.api.OpenMaya as om
+    import maya.cmds as cmds
+    import maya.mel as mel
+except ImportError:
+    # Backward compatibility to support test in uinstalled maya.
+    try:
+        # Import built-in modules
+        from unittest.mock import MagicMock
+    except ImportError:
+        # Import third-party modules
+        from mock import MagicMock  # noqa: UP026
+    cmds = MagicMock()
+    om = MagicMock()
+    mel = MagicMock()
+
+# Import built-in modules
+from functools import wraps
+
+
+def is_maya_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 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"^\['.+']")

maya_umbrella/vaccine.py

@@ -1,207 +1,12 @@
-# Import built-in modules
-from collections import defaultdict
-import glob
-import logging
-import os
-
-# Import third-party modules
-import maya.cmds as cmds
-
-# Import local modules
-from maya_umbrella.filesystem import safe_remove_file
-from maya_umbrella.filesystem import safe_rmtree
-
-
-class MayaVirusCleaner(object):
-    _bad_files = []
-    _bad_nodes = []
-    _bad_script_nodes = []
-    _bad_script_jobs = []
-    _registered_callbacks = defaultdict(list)
-    _fix_funcs = []
-
-    def __init__(self, logger=None, auto_fix=True):
-        self.logger = logger or logging.getLogger(__name__)
-        self.auto_fix = auto_fix
-
-    @property
-    def user_app_dir(self):
-        """Return the user application directory."""
-        return cmds.internalVar(userAppDir=True)
-
-    @property
-    def maya_file(self):
-        """Return the current Maya file."""
-        return cmds.file(q=True, sn=True)
-
-    @property
-    def maya_install_root(self):
-        """Return the Maya installation root directory."""
-        return os.environ["MAYA_LOCATION"]
-
-    @property
-    def user_script_path(self):
-        """Return the user script directory."""
-        return cmds.internalVar(userScriptDir=True)
-
-    @property
-    def local_script_path(self):
-        """Return the local script directory."""
-        return os.path.join(self.user_app_dir, "scripts")
-
-    @property
-    def bad_files(self):
-        """Return a list of bad files."""
-        return [path for path in list(set(self._bad_files)) if os.path.exists(path)]
-
-    @property
-    def bad_nodes(self):
-        """Return a list of bad nodes."""
-        return list(set(self._bad_nodes))
-
-    @property
-    def bad_script_nodes(self):
-        """Return a list of bad script nodes."""
-        return list(set(self._bad_script_nodes))
-
-    @property
-    def bad_script_jobs(self):
-        """Return a list of bad script jobs."""
-        return list(set(self._bad_script_jobs))
-
-    def callback_remove_rename_temp_files(self, *args, **kwargs):
-        """
-        Remove temporary files in the local script path.
-        """
-        self.logger.info("Removing temporary files in %s", self.local_script_path)
-        [safe_remove_file(temp_file) for temp_file in glob.glob(os.path.join(self.local_script_path, "._*"))]
-
-    @property
-    def registered_callbacks(self):
-        return self._registered_callbacks
-
-    def add_bad_files(self, files):
-        self._bad_files.extend(files)
-
-    def add_bad_file(self, file):
-        self._bad_files.append(file)
-
-    def add_bad_nodes(self, nodes):
-        self._bad_nodes.extend(nodes)
-
-    def add_bad_node(self, node):
-        self._bad_nodes.append(node)
-
-    def add_bad_script_jobs(self, jobs):
-        self._bad_script_jobs.extend(jobs)
-
-    def add_bad_script_job(self, job):
-        self._bad_script_jobs.append(job)
-
-    def add_bad_script_nodes(self, nodes):
-        self._bad_script_nodes.extend(nodes)
-
-    def add_bad_script_node(self, node):
-        self._bad_script_nodes.append(node)
-
-    def register_callback(self, callback_name, callback):
-        """Register a callback to be executed before or after processing."""
-        self._registered_callbacks[callback_name].append(callback)
-
-    def add_after_open_callback(self, callback):
-        self.register_callback("after_open", callback)
-
-    def add_maya_initialized_callback(self, callback):
-        self.register_callback("maya_initialized", callback)
-
-    def add_after_import_callback(self, callback):
-        self.register_callback("after_import", callback)
-
-    def add_after_import_reference_callback(self, callback):
-        self.register_callback("after_import_reference", callback)
-
-    def add_after_load_reference_callback(self, callback):
-        self.register_callback("after_load_reference", callback)
-
-    def add_before_save_callback(self, callback):
-        self.register_callback("before_save", callback)
-
-    def add_before_import_callback(self, callback):
-        self.register_callback("before_import", callback)
-
-    def add_before_load_reference_callback(self, callback):
-        self.register_callback("before_load_reference", callback)
-
-    def add_before_import_reference_callback(self, callback):
-        self.register_callback("before_import_reference", callback)
-
-    def add_maya_exiting_callback(self, callback):
-        self.register_callback("maya_exiting", callback)
-
-    def setup_default_callbacks(self):
-        self.add_maya_initialized_callback(self.callback_remove_rename_temp_files)
-        self.add_maya_exiting_callback(self.callback_remove_rename_temp_files)
-
-    def add_fix_function(self, func):
-        self._fix_funcs.append(func)
-
-    # fix
-
-    def fix_script_jobs(self):
-        for script_job in self.bad_script_jobs:
-            script_num = int(script_job.split(":", 1)[0])
-            self.logger.info("Kill script job %s", script_job)
-            cmds.scriptJob(kill=script_num, force=True)
-            self._bad_script_jobs.remove(script_job)
-
-    def fix_bad_files(self):
-        for file_ in self.bad_files:
-            if os.path.exists(file_):
-                if os.path.isfile(file_):
-                    self.logger.info("Removing %s", file_)
-                    safe_remove_file(file_)
-                    self._bad_files.remove(file_)
-                else:
-                    self.logger.info("Removing folder %s", file_)
-                    safe_rmtree(file_)
-                    self._bad_files.remove(file_)
-
-    def fix_bad_nodes(self):
-        for node in self.bad_nodes:
-            self.logger.info("Deleting %s", node)
-            try:
-                cmds.lockNode(node, lock=False)
-            except ValueError:
-                pass
-            try:
-                cmds.delete(node)
-            except ValueError:
-                pass
-            self._bad_nodes.remove(node)
-
-    def fix_all_issues(self):
-        """Fix all issues related to the Maya virus."""
-        self.fix_bad_files()
-        self.fix_bad_nodes()
-        self.fix_script_jobs()
-        for func in self._fix_funcs:
-            func()
-
-    def report_all_issues(self):
-        """Report all issues related to the Maya virus."""
-        self.logger.info("Bad files: %s", self.bad_files)
-        self.logger.info("Bad nodes: %s", self.bad_nodes)
-        self.logger.info("Bad script jobs: %s", self.bad_script_jobs)
-
-    def reset_all_issues(self):
-        """Reset all issues related to the Maya virus."""
-        self._bad_files = []
-        self._bad_nodes = []
-        self._bad_script_nodes = []
-        self._bad_script_jobs = []
+class AbstractVaccine(object):
+    """Abstract base class for Vaccine classes.
 
+    Attributes:
+        virus_name (str): The name of the virus.
+        api (MayaVirusCleaner): The VaccineAPI instance.
+        logger (Logger): The logger instance.
+    """
 
-class AbstractVaccine(object):
     virus_name = None
 
     def __init__(self, api, logger):
@@ -216,7 +21,17 @@ class AbstractVaccine(object):
         self.logger = logger
 
     def collect_issues(self):
+        """Collect issues related to the virus.
+
+        Raises:
+            NotImplementedError: This method must be implemented in the derived classes.
+        """
         raise NotImplementedError
 
     def report_issue(self, name):
-        self.logger.warning("%s: Infected by Malware!", name)
+        """Report an issue related to the virus.
+
+        Args:
+            name (str): The name of the issue.
+        """
+        self.logger.warning(self.api.translator.translate("report_issue", name=name))

maya_umbrella/vaccines/vaccine1.py

@@ -6,10 +6,11 @@ from maya_umbrella.vaccine import AbstractVaccine
 
 
 class Vaccine(AbstractVaccine):
-    virus_name = "pu tian tong qi"
+    """A class for handling the PuTianTongQi virus."""
+    virus_name = "PutTianTongQi"
 
     def collect_issues(self):
-        self.api.add_bad_files(
+        self.api.add_malicious_files(
             [
                 os.path.join(self.api.local_script_path, "fuckVirus.py"),
                 os.path.join(self.api.local_script_path, "fuckVirus.pyc"),

maya_umbrella/vaccines/vaccine2.py

@@ -1,60 +1,52 @@
 # Import built-in modules
 import os.path
 
-# Import third-party modules
-import maya.cmds as cmds
-
 # Import local modules
-from maya_umbrella.filesystem import read_file
-from maya_umbrella.filesystem import rename
+from maya_umbrella.constants import JOB_SCRIPTS_VIRUS_SIGNATURES
+from maya_umbrella.filesystem import check_virus_by_signature
+from maya_umbrella.filesystem import check_virus_file_by_signature
+from maya_umbrella.maya_funs import check_reference_node_exists
+from maya_umbrella.maya_funs import cmds
+from maya_umbrella.maya_funs import get_attr_value
 from maya_umbrella.vaccine import AbstractVaccine
 
 
 class Vaccine(AbstractVaccine):
+    """A class for handling the ZeiJianKang virus."""
+
     virus_name = "zei jian kang"
 
-    def collect_bad_nodes(self):
+    def collect_infected_nodes(self):
         """Collect all bad nodes related to the virus."""
         for script_node in cmds.ls(type="script"):
-            if cmds.referenceQuery(script_node, isNodeReferenced=True):
+            if check_reference_node_exists(script_node):
                 continue
-            script_before_string = cmds.getAttr("{}.before".format(script_node))
-            script_after_string = cmds.getAttr("{}.after".format(script_node))
-            for script_string in [script_before_string, script_after_string]:
+            for attr_name in ("before", "after"):
+                script_string = get_attr_value(script_node, attr_name)
                 if not script_string:
                     continue
-                if "internalVar" in script_string or "userSetup" in script_string or "fuckVirus" in script_string:
+                if check_virus_by_signature(script_string, JOB_SCRIPTS_VIRUS_SIGNATURES):
                     self.report_issue(script_node)
-                    self.api.add_bad_node(script_node)
+                    self.api.add_infected_node(script_node)
 
     def collect_issues(self):
         """Collect all issues related to the virus."""
-        self.api.add_bad_files(
+        self.api.add_malicious_files(
             [
                 os.path.join(self.api.local_script_path, "vaccine.py"),
                 os.path.join(self.api.local_script_path, "vaccine.pyc"),
             ],
         )
-        self.collect_bad_usersetup_py()
-        self.collect_bad_nodes()
+        self.collect_infected_user_setup_py()
+        self.collect_infected_nodes()
 
-    def collect_bad_usersetup_py(self):
+    def collect_infected_user_setup_py(self):
         """Collect all bad userSetup.py files related to the virus."""
-        for usersetup_py in [
-            os.path.join(self.api.local_script_path, "vaccine.py"),
-            os.path.join(self.api.user_script_path, "vaccine.py"),
+        for user_setup_py in [
             os.path.join(self.api.local_script_path, "userSetup.py"),
             os.path.join(self.api.user_script_path, "userSetup.py"),
         ]:
-            if os.path.exists(usersetup_py):
-                data = read_file(usersetup_py)
-                if "petri_dish_path = cmds.internalVar(userAppDir=True) + 'scripts/userSetup.py" in data:
-                    self.report_issue("vaccine.py")
-                    self.api.add_bad_file(rename(usersetup_py))
-
-                if (
-                    "cmds.evalDeferred('leukocyte = vaccine.phage()')" in data
-                    and "cmds.evalDeferred('leukocyte.occupation()')" in data
-                ):
-                    self.report_issue("userSetup.py")
-                    self.api.add_bad_file(rename(usersetup_py))
+            if os.path.exists(user_setup_py):
+                if check_virus_file_by_signature(user_setup_py):
+                    self.report_issue(user_setup_py)
+                    self.api.add_infected_file(user_setup_py)