feat: Refactor code logic to make API more flexible.

Signed-off-by: loonghao <[email protected]>

.coveragerc

@@ -11,3 +11,4 @@ exclude_lines =
 ignore_errors = True
 omit =
     tests/*
+    maya_umbrella/maya_funs.py

.github/workflows/mr-test.yml

@@ -27,3 +27,6 @@ jobs:
       - name: isort
         run: |
           nox -s isort_check
+      - name: make-zip
+        run: |
+          nox -s make-zip -- --version 0.5.0

README.md

@@ -27,6 +27,12 @@ It ensures a secure and seamless user experience by proactively scanning for thr
 
 It can be provided as an API for seamless integration into your existing pipeline.
 
+# 安装
+maya_umbrella是以标准pipy包去发布的，所以我们可以通过pip install去安装
+```shell
+your/maya-root/mayapy -m pip install maya-umbrella
+```
+
 # 开发环境设置
 
 Set up the development environment using a virtual environment,
@@ -78,16 +84,6 @@ nox -s ruff_check
 nox -s make-zip -- --version 0.5.0
 ```
 
-# API
-获取当前场景没有被修复的病毒文件
-
-```python
-from maya_umbrella import MayaVirusDefender
-
-api = MayaVirusDefender()
-print(api.get_unfixed_references)
-```
-
 # 环境变量
 
 我们可以通过下列环境变量去修改maya_umbrella的一些设置，方便有pipeline的公司可以更好的集成
@@ -108,6 +104,45 @@ MAYA_UMBRELLA_LOG_NAME
 
 ```shell
 MAYA_UMBRELLA_LOG_LEVEL
+```
+修改杀毒后文件的备份文件夹名称， 默认是`_maya_umbrella`
+比如:
+你文件路径是  `c:/your/path/file.ma`
+那么备份文件路径是 `c:/your/path/_maya_umbrella/file.ma`
+```shell
+MAYA_UMBRELLA_BACKUP_FOLDER_NAME
+```
+默认的显示语言，包含日志打印输出等，默认是根据你当前maya的界面语言来设置的，当然我们也可以通过下面的环境变量去设置
+```shell
+MAYA_UMBRELLA_LANG
+```
+
+忽略保存到备份文件夹，*请注意，如果你不清楚这个会导致的后果请不要轻易修改*，默认批量杀毒后会自动保存到当前文件的备份文件夹.
+```shell
+MAYA_UMBRELLA_IGNORE_BACKUP
+```
+如果忽略请设置为
+```shell
+SET MAYA_UMBRELLA_IGNORE_BACKUP=true
+```
+
+# API
+获取当前场景没有被修复的病毒文件
+
+```python
+from maya_umbrella import MayaVirusDefender
+
+api = MayaVirusDefender()
+print(api.get_unfixed_references())
+```
+
+批量修复文件
+```python
+from maya_umbrella import MayaVirusScanner
+
+api = MayaVirusScanner()
+print(api.scan_files_from_pattern("your/path/*.m[ab]"))
+
 ```
 
 ## Contributors ✨

dev/get-pip-2.7.py

@@ -6812,7 +6812,7 @@ def main():
 )q(qUQNG<5=<AK)z>N_AG@V%iG%#8oeCF6gbXHk%b8P6gu^oOo^c%&wL_DK!uX}+MyYmqm@FbgV8yTj
 Q(bFTQ%OnardAP6oYIPM5SURU0?74~YqPN>Cb$LXGOBke&L#nw~nTFE#*bVJ+@5$L!jnuo`s|Po4;Qc
 M#7WoEYZ(~$<rS!7ee$f5PD}4r8dNjB;$h@k4Ku@bq3&HA#eXj9pAx4K%zfhMvGbX~yc>Eam*ZJx6E{
-0uOz7gB(X4T^DsLBUZIW7srGhVQqZx3~ByO`L^H}loc+R)+$0X`D^(aK(<ejQNRdq@C0XWu+2p1RGe@
+0uOz7gB(X4T^DsLBUZIW7srGhVQqZx3~ByO`L^H}locale+R)+$0X`D^(aK(<ejQNRdq@C0XWu+2p1RGe@
 ov`mJ%pCWrbDvy!H4&JNS453nEBzdmY1{KY<J|Dv>x23UD=l8=lsk!IL))8$g<&yy9qo8HP#0)+8gv6
 j=;kNz6F02R=Q%s){0CVYXkZKAZ%fikG8ld<%ZpCls)92alG%j2lpg`r`kR-W(=+cKI-{<fQ;km@o2F
 fxc*$l{W{#6a4fXL5!L62CvMQ!r)^oEwqrxo3D3;ydvNrW`PhI)0-L>Oy=bB=ux<`s6Fg76U>eg=Ze<

manual_test_in_maya.py

@@ -1,38 +1,16 @@
 import os
 import glob
 import maya.cmds as cmds
-from functools import wraps
+from maya_umbrella.maya_funs import open_maya_file
 
 ROOT = os.path.dirname(os.path.abspath(__file__))
 
 
-def block_prompt(func):
-    @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
-
-
 def get_virus_files():
     return glob.glob(os.path.join(ROOT, "tests", "virus", "*.ma"))
 
 
-@block_prompt
 def start():
     for maya_file in get_virus_files():
-        cmds.file(new=True, force=True)
-        try:
-            cmds.file(maya_file, open=True, force=True, ignoreVersion=True, executeScriptNodes=False)
-        except Exception:
-            pass
+        open_maya_file(maya_file)
         cmds.file(new=True, force=True)

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/cleaner.py

@@ -0,0 +1,110 @@
+# 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.info("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.info("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.info("Removing %s", file_)
+                    safe_remove_file(file_)
+                    self.collector.remove_malicious_file(file_)
+                else:
+                    self.logger.info("Removing folder %s", 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:
+            self.logger.info(self.translator.translate("delete", name=node))
+            is_referenced = check_reference_node_exists(node)
+            if is_referenced:
+                try:
+                    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:
+                    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("Starting Fixing all issues related to the Maya virus from %s.", 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("Finished Fixing all issues related to the Maya virus from %s.", maya_file)