maya-umbrella 0.5.0__py2.py3-none-any.whl → 0.6.1__py2.py3-none-any.whl

maya_umbrella/__init__.py

@@ -1,5 +1,9 @@
 # Import local modules
-from maya_umbrella.core import MayaVirusDefender
+from maya_umbrella.cleaner import MayaVirusCleaner
+from maya_umbrella.collector import MayaVirusCollector
+from maya_umbrella.defender import MayaVirusDefender
+from maya_umbrella.defender import context_defender
+from maya_umbrella.scanner import MayaVirusScanner
 
 
-__all__ = ["MayaVirusDefender"]
+__all__ = ["MayaVirusDefender", "MayaVirusCleaner", "MayaVirusCollector", "MayaVirusScanner", "context_defender"]

maya_umbrella/__version__.py

@@ -1 +1 @@
-__version__ = "0.5.0"
+__version__ = "0.6.1"

maya_umbrella/cleaner.py

@@ -0,0 +1,111 @@
+# Import built-in modules
+import glob
+import logging
+import os
+import re
+
+# Import local modules
+from maya_umbrella.constants import FILE_VIRUS_SIGNATURES
+from maya_umbrella.filesystem import remove_virus_file_by_signature
+from maya_umbrella.filesystem import safe_remove_file
+from maya_umbrella.filesystem import safe_rmtree
+from maya_umbrella.i18n import Translator
+from maya_umbrella.maya_funs import check_reference_node_exists
+from maya_umbrella.maya_funs import cmds
+
+
+class MayaVirusCleaner(object):
+    """A class to clean Maya virus files.
+
+    Attributes:
+        logger (Logger): Logger object for logging purposes.
+        translator (Translator): Translator object for translation purposes.
+        collector (MayaVirusCollector): MayaVirusCollector object for collecting issues.
+    """
+    def __init__(self, collector, logger=None):
+        """Initialize the MayaVirusCleaner.
+
+        Args:
+            collector (MayaVirusCollector): MayaVirusCollector object for collecting issues.
+            logger (Logger, optional): Logger object for logging purposes. Defaults to None, which creates a new logger.
+        """
+        self.logger = logger or logging.getLogger(__name__)
+        self.translator = Translator()
+        self.collector = collector
+
+    def callback_remove_rename_temp_files(self, *args, **kwargs):
+        """Remove temporary files in the local script path."""
+        self.logger.debug("Removing temporary files in %s", self.collector.local_script_path)
+        for temp_file in glob.glob(os.path.join(self.collector.local_script_path, "._*")):
+            safe_remove_file(temp_file)
+
+    def fix_script_jobs(self):
+        """Fix infected script jobs."""
+        for script_job in self.collector.infected_script_jobs:
+            script_num = int(re.findall(r"^(\d+):", script_job)[0])
+            self.logger.debug("Kill script job %s", script_job)
+            cmds.scriptJob(kill=script_num, force=True)
+            self.collector.remove_infected_script_job(script_job)
+
+    def fix_malicious_files(self):
+        """Fix malicious files."""
+        for file_ in self.collector.malicious_files:
+            if os.path.exists(file_):
+                if os.path.isfile(file_):
+                    self.logger.debug(self.translator.translate("remove_file", name=file_))
+                    safe_remove_file(file_)
+                    self.collector.remove_malicious_file(file_)
+                else:
+                    self.logger.debug(self.translator.translate("remove_path", name=file_))
+                    safe_rmtree(file_)
+                    self.collector.remove_malicious_file(file_)
+
+    def fix_infected_nodes(self):
+        """Fix infected nodes."""
+        for node in self.collector.infected_nodes:
+            is_referenced = check_reference_node_exists(node)
+            if is_referenced:
+                try:
+                    self.logger.debug(self.translator.translate("fix_infected_reference_nodes", name=node))
+                    cmds.setAttr("{node}.before".format(node=node), "", type="string")
+                    cmds.setAttr("{node}.after".format(node=node), "", type="string")
+                    cmds.setAttr("{node}.scriptType".format(node=node), 0)
+                    self.collector.remove_infected_node(node)
+                except Exception as e:
+                    self.logger.debug(e)
+            else:
+                try:
+                    cmds.lockNode(node, lock=False)
+                except ValueError:
+                    pass
+                try:
+                    self.logger.debug(self.translator.translate("fix_infected_nodes", name=node))
+                    cmds.delete(node)
+                except ValueError:
+                    pass
+                self.collector.remove_infected_node(node)
+
+    def setup_default_callbacks(self):
+        """Set up default callbacks."""
+        self.collector.add_maya_initialized_callback(self.callback_remove_rename_temp_files)
+        self.collector.add_maya_exiting_callback(self.callback_remove_rename_temp_files)
+
+    def fix_infected_files(self):
+        """Fix infected files."""
+        for file_path in self.collector.infected_files:
+            self.logger.info(self.translator.translate("fix_infected_files", name=file_path))
+            remove_virus_file_by_signature(file_path, FILE_VIRUS_SIGNATURES)
+            self.collector.remove_infected_file(file_path)
+
+    def fix(self):
+        """Fix all issues related to the Maya virus."""
+        if self.collector.have_issues:
+            maya_file = cmds.file(query=True, sceneName=True, shortName=True) or "empty/scene"
+            self.logger.info(self.translator.translate("start_fix_issues", name=maya_file))
+            self.fix_malicious_files()
+            self.fix_infected_files()
+            self.fix_infected_nodes()
+            self.fix_script_jobs()
+            for func in self.collector.get_additionally_fix_funcs():
+                func()
+            self.logger.info(self.translator.translate("finish_fix_issues", name=maya_file))

maya_umbrella/collector.py

@@ -0,0 +1,411 @@
+# Import built-in modules
+from collections import defaultdict
+import logging
+import os
+
+# Import local modules
+from maya_umbrella.filesystem import get_vaccines
+from maya_umbrella.filesystem import load_hook
+from maya_umbrella.i18n import Translator
+from maya_umbrella.maya_funs import cmds
+
+
+class MayaVirusCollector(object):
+    """A class to collect and handle Maya viruses.
+
+    Attributes:
+        _malicious_files (list): List to store malicious files.
+        _infected_files (list): List to store infected files.
+        _infected_nodes (list): List to store infected nodes.
+        _infected_script_nodes (list): List to store infected script nodes.
+        _infected_reference_files (list): List to store infected reference files.
+        _infected_script_jobs (list): List to store infected script jobs.
+        _registered_callbacks (defaultdict): Dictionary to store registered callbacks.
+        _additionally_fix_funcs (list): List to store additional fix functions.
+        _vaccines (list): List to store vaccines.
+        logger: Logger object for logging purposes.
+        translator: Translator object for translation purposes.
+    """
+
+    def __init__(self, logger, translator=None):
+        """Initialize MayaVirusCollector.
+
+        Args:
+            logger (Logger, optional): Logger object for logging purposes.
+            translator (Translator, optional): Translator object for translation purposes.
+        """
+        self.logger = logger or logging.getLogger(__name__)
+        self.translator = translator or Translator()
+        # Malicious files or temp files that need to be deleted directly.
+        self._malicious_files = []
+        self._infected_files = []
+        self._infected_nodes = []
+        self._infected_script_nodes = []
+        self._infected_reference_files = []
+        self._infected_script_jobs = []
+        self._registered_callbacks = defaultdict(list)
+        self._additionally_fix_funcs = []
+        self._vaccines = []
+        self.load_vaccines()
+
+    def load_vaccines(self):
+        """Load all vaccines."""
+        for vaccine in get_vaccines():
+            vaccine_class = load_hook(vaccine).Vaccine
+            try:
+                self._vaccines.append(vaccine_class(api=self, logger=self.logger))
+            except Exception as e:
+                self.logger.error("Error loading vaccine: %s", e)
+
+    @property
+    def vaccines(self):
+        """Get all loaded vaccines.
+
+        Returns:
+            list: A list of loaded vaccines.
+        """
+        return self._vaccines
+
+    @property
+    def user_app_dir(self):
+        """Return the user application directory."""
+        return cmds.internalVar(userAppDir=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 malicious_files(self):
+        """Return a list of bad files."""
+        return [path for path in list(set(self._malicious_files)) if os.path.exists(path)]
+
+    @property
+    def infected_nodes(self):
+        """Return a list of infected nodes."""
+        return list(set(self._infected_nodes))
+
+    @property
+    def infected_reference_files(self):
+        return [path for path in list(set(self._infected_reference_files)) if os.path.exists(path)]
+
+    @property
+    def infected_script_nodes(self):
+        """Return a list of bad script nodes."""
+        return list(set(self._infected_script_nodes))
+
+    @property
+    def infected_script_jobs(self):
+        """Return a list of bad script jobs."""
+        return list(set(self._infected_script_jobs))
+
+    @property
+    def infected_files(self):
+        """Return the list of infected files.
+
+        Returns:
+            list: List of infected files.
+        """
+        return self._infected_files
+
+    @property
+    def registered_callbacks(self):
+        """Return the dictionary of registered callbacks.
+
+        Returns:
+            defaultdict: Dictionary of registered callbacks.
+        """
+        return self._registered_callbacks
+
+    def add_infected_reference_files(self, files):
+        """Add multiple infected reference files.
+
+        Args:
+            files (list): List of infected reference files to be added.
+        """
+        self._infected_reference_files.extend(files)
+
+    def add_infected_reference_file(self, file):
+        """Add a single infected reference file.
+
+        Args:
+            file (str): Infected reference file to be added.
+        """
+        self._infected_reference_files.append(file)
+
+    def remove_infected_reference_file(self, file):
+        """Remove an infected reference file.
+
+        Args:
+            file (str): Infected reference file to be removed.
+        """
+        self._infected_reference_files.remove(file)
+
+    def add_infected_files(self, files):
+        """Add multiple infected files.
+
+        Args:
+            files (list): List of infected files to be added.
+        """
+        self._infected_files.extend(files)
+
+    def add_infected_file(self, file):
+        """Add a single infected file.
+
+        Args:
+            file (str): Infected file to be added.
+        """
+        self._infected_files.append(file)
+
+    def remove_infected_file(self, file):
+        """Remove an infected file.
+
+        Args:
+            file (str): Infected file to be removed.
+        """
+        self._infected_files.remove(file)
+
+    def add_malicious_files(self, files):
+        """Add multiple malicious files.
+
+        Args:
+            files (list): List of malicious files to be added.
+        """
+        self._malicious_files.extend(files)
+
+    def add_malicious_file(self, file):
+        """Add a single malicious file.
+
+        Args:
+            file (str): Malicious file to be added.
+        """
+        self._malicious_files.append(file)
+
+    def remove_malicious_file(self, file):
+        """Remove a malicious file.
+
+        Args:
+            file (str): Malicious file to be removed.
+        """
+        self._malicious_files.remove(file)
+
+    def add_infected_nodes(self, nodes):
+        """Add multiple infected nodes.
+
+        Args:
+            nodes (list): List of infected nodes to be added.
+        """
+        self._infected_nodes.extend(nodes)
+
+    def add_infected_node(self, node):
+        """Add a single infected node.
+
+        Args:
+            node (str): Infected node to be added.
+        """
+        self._infected_nodes.append(node)
+
+    def remove_infected_node(self, node):
+        """Remove an infected node.
+
+        Args:
+            node (str): Infected node to be removed.
+        """
+        self._infected_nodes.remove(node)
+
+    def add_infected_script_jobs(self, jobs):
+        """Add multiple infected script jobs.
+
+        Args:
+            jobs (list): List of infected script jobs to be added.
+        """
+        self._infected_script_jobs.extend(jobs)
+
+    def add_infected_script_job(self, job):
+        """Add a single infected script job.
+
+        Args:
+            job (str): Infected script job to be added.
+        """
+        self._infected_script_jobs.append(job)
+
+    def remove_infected_script_job(self, job):
+        """Remove an infected script job.
+
+        Args:
+            job (str): Infected script job to be removed.
+        """
+        self._infected_script_jobs.remove(job)
+
+    def add_infected_script_nodes(self, nodes):
+        """Add multiple infected script nodes.
+
+        Args:
+            nodes (list): List of infected script nodes to be added.
+        """
+        self._infected_script_nodes.extend(nodes)
+
+    def add_infected_script_node(self, node):
+        """Add a single infected script node.
+
+        Args:
+            node (str): Infected script node to be added.
+        """
+        self._infected_script_nodes.append(node)
+
+    def remove_infected_script_node(self, node):
+        """Remove an infected script node.
+
+        Args:
+            node (str): Infected script node to be removed.
+        """
+        self._infected_script_nodes.remove(node)
+
+    def get_additionally_fix_funcs(self):
+        """Return the list of additional fix functions.
+
+        Returns:
+            list: List of additional fix functions.
+        """
+        return self._additionally_fix_funcs
+
+    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):
+        """Add a callback to be executed after a file is opened.
+
+        Args:
+            callback (function): The callback function to be added.
+        """
+        self.register_callback("after_open", callback)
+
+    def add_maya_initialized_callback(self, callback):
+        """Add a callback to be executed when Maya is initialized.
+
+        Args:
+            callback (function): The callback function to be added.
+        """
+        self.register_callback("maya_initialized", callback)
+
+    def add_after_import_callback(self, callback):
+        """Add a callback to be executed after a file is imported.
+
+        Args:
+            callback (function): The callback function to be added.
+        """
+        self.register_callback("after_import", callback)
+
+    def add_after_import_reference_callback(self, callback):
+        """Add a callback to be executed after a reference file is imported.
+
+        Args:
+            callback (function): The callback function to be added.
+        """
+        self.register_callback("after_import_reference", callback)
+
+    def add_after_load_reference_callback(self, callback):
+        """Add a callback to be executed after a reference file is loaded.
+
+        Args:
+            callback (function): The callback function to be added.
+        """
+        self.register_callback("after_load_reference", callback)
+
+    def add_before_save_callback(self, callback):
+        """Add a callback to be executed before a file is saved.
+
+        Args:
+            callback (function): The callback function to be added.
+        """
+        self.register_callback("before_save", callback)
+
+    def add_before_import_callback(self, callback):
+        """Add a callback to be executed before a file is imported.
+
+        Args:
+            callback (function): The callback function to be added.
+        """
+        self.register_callback("before_import", callback)
+
+    def add_before_load_reference_callback(self, callback):
+        """Add a callback to be executed before a reference file is loaded.
+
+        Args:
+            callback (function): The callback function to be added.
+        """
+        self.register_callback("before_load_reference", callback)
+
+    def add_before_import_reference_callback(self, callback):
+        """Add a callback to be executed before a reference file is imported.
+
+        Args:
+            callback (function): The callback function to be added.
+        """
+        self.register_callback("before_import_reference", callback)
+
+    def add_maya_exiting_callback(self, callback):
+        """Add a callback to be executed when Maya is exiting.
+
+        Args:
+            callback (function): The callback function to be added.
+        """
+        self.register_callback("maya_exiting", callback)
+
+    def add_additionally_fix_function(self, func):
+        """Add an additional fix function to be executed.
+
+        Args:
+            func (function): The fix function to be added.
+        """
+        self._additionally_fix_funcs.append(func)
+
+    def collect(self):
+        """Collect issues from all loaded vaccines."""
+        self.reset()
+        for vaccine in self.vaccines:
+            vaccine.collect_issues()
+
+    @property
+    def have_issues(self):
+        """Check if any issues are found.
+
+        Returns:
+            bool: True if any issues are found, False otherwise.
+        """
+        return any([self.malicious_files, self.infected_nodes, self.infected_script_nodes, self.infected_files,
+                    self.infected_script_jobs])
+
+    def report(self):
+        """Report all issues related to the Maya virus."""
+        for name in (
+            "malicious_files",
+            "infected_nodes",
+            "infected_script_jobs",
+            "infected_files",
+            "infected_reference_files"
+        ):
+            self.logger.info(self.translator.translate(name, name=getattr(self, name)))
+
+    def reset(self):
+        """Reset all issues related to the Maya virus."""
+        self._malicious_files = []
+        self._infected_nodes = []
+        self._infected_script_nodes = []
+        self._infected_script_jobs = []
+        self._infected_files = []
+        self._infected_reference_files = []
+        self._registered_callbacks = defaultdict(list)
+        self._additionally_fix_funcs = []

maya_umbrella/constants.py

@@ -7,6 +7,7 @@ LOG_MAX_BYTES = 1024 * 1024 * 5
 FILE_VIRUS_SIGNATURES = [
     "import vaccine",
     "cmds.evalDeferred.*leukocyte.+",
+    # https://regex101.com/r/0MNzF7/1
     "python(.*);.+exec.+(pyCode).+;",
 ]
 
@@ -14,5 +15,8 @@ JOB_SCRIPTS_VIRUS_SIGNATURES = [
     "petri_dish_path.+cmds.internalVar.+",
     "userSetup",
     "fuckVirus",
+    # https://regex101.com/r/0MNzF7/1
     "python(.*);.+exec.+(pyCode).+;",
+    # https://regex101.com/r/2D14UA/1
+    r"^\['.+']",
 ]