Merge pull request #6 from cpcgskill/develop

feat: ✨ 检查支持mb文件，支持清除mb文件的病毒

src/gui.py

@@ -52,7 +52,7 @@ def clear_virus_script_node():
         cmds.warning('未发现病毒脚本节点')
 
 
-def batch_clear_maya_file():
+def batch_clear_maya_file_ma():
     """
     批量清理Maya文件.
     """
@@ -77,7 +77,42 @@ def batch_clear_maya_file():
         return
     for file in files:
         try:
-            kill_mayafile.process_file(os.path.abspath(file))
+            kill_mayafile.process_ma_file(os.path.abspath(file))
+        except Exception as e:
+            import traceback
+            traceback.print_exc()
+            cmds.warning('处理文件{}失败: {}'.format(file, e))
+        else:
+            cmds.warning('处理文件{}成功'.format(file))
+    cmds.warning('批量处理完成')
+
+
+def batch_clear_maya_file_mb():
+    """
+    批量清理Maya文件.
+    """
+    # 弹窗警告
+    result = cmds.confirmDialog(
+        title='警告',
+        message='正则批量清理Maya文件, 可能会导致文件损坏, 请注意备份',
+        button=['Ok', 'Cancel'],
+        defaultButton='Cancel',
+        cancelButton='Cancel',
+        dismissString='Cancel',
+    )
+    if result == 'Cancel':
+        cmds.warning('已取消')
+        return
+
+    #
+    cmds.warning('开始批量处理...')
+    files = cmds.fileDialog2(fileMode=4, dialogStyle=2)
+    if not files:
+        cmds.warning('未选择文件')
+        return
+    for file in files:
+        try:
+            kill_mayafile.process_mb_file(os.path.abspath(file))
         except Exception as e:
             import traceback
             traceback.print_exc()
@@ -114,7 +149,13 @@ def batch_check_maya_file():
     for root, dirs, files in os.walk(os.path.abspath(root_dir[0])):
         for file in files:
             if file.endswith('.ma'):
-                has_virus = kill_mayafile.check_file(os.path.join(root, file), is_cautious=is_cautious)
+                cmds.warning('检查文件: {}'.format(os.path.join(root, file)))
+                has_virus = kill_mayafile.check_ma_file(os.path.join(root, file), is_cautious=is_cautious)
+                if has_virus:
+                    warning_files.append(os.path.join(root, file).replace('\\', '/'))
+            if file.endswith('.mb'):
+                cmds.warning('检查文件: {}'.format(os.path.join(root, file)))
+                has_virus = kill_mayafile.check_mb_file(os.path.join(root, file), is_cautious=is_cautious)
                 if has_virus:
                     warning_files.append(os.path.join(root, file).replace('\\', '/'))
 
@@ -146,8 +187,9 @@ def create_gui():
     cmds.button(label='单独清除HIK病毒', command=lambda *args: alone_clear_hik_virus())
     cmds.button(label='恢复UAC设置', command=lambda *args: restore_UAC())
     cmds.button(label='清除病毒脚本节点', command=lambda *args: clear_virus_script_node())
-    cmds.button(label='批量清理Maya文件(.ma)', command=lambda *args: batch_clear_maya_file())
-    cmds.button(label='批量检查Maya文件(.ma)', command=lambda *args: batch_check_maya_file())
+    cmds.button(label='批量清理Maya文件(.ma)', command=lambda *args: batch_clear_maya_file_ma())
+    cmds.button(label='批量清理Maya文件(.mb)', command=lambda *args: batch_clear_maya_file_mb())
+    cmds.button(label='批量检查Maya文件', command=lambda *args: batch_check_maya_file())
 
     cmds.showWindow('VirusKiller_240429')
 

src/kill_mayafile.py

@@ -15,56 +15,83 @@
 import sys
 import maya.cmds as cmds
 
-find_script_node_regex = re.compile(
-    r'\s*createNode\s+script\s+-n\s+\"uifiguration\";\s*.*?setAttr\s*\"\.b\"\s*-type\s*"string".*?\s*setAttr\s*"\.st"\s*1;\s*',
+ma_find_script_node_regex = re.compile(
+    r'\s*createNode\s+script\s+-n\s+\"uifiguration\";\s*.*?setAttr\s*\"\.b\"\s*-type\s*"string".*?\s*setAttr\s*"\.st"\s*1;\s*'.encode(),
     re.DOTALL,
 )
+cautious_check_regex = re.compile(
+    r'\s*import\s*base64;?\s*'.encode(),
+    re.DOTALL,
+)
+mb_regex = re.compile(
+    r'\s*?python\s*?\(\s*?\"\s*?import\s+?base64;.*?;\s*?exec\s*?\(\s*?_pycode\s*?\)\s*?\"\s*?\)\s*?;\s*?'.encode(),
+    re.DOTALL
+)
 
 
-def _process_str(data):
-    for i in reversed(list(find_script_node_regex.finditer(data))):
+def _process_ma_content(data):
+    for i in reversed(list(ma_find_script_node_regex.finditer(data))):
         start, end = i.span()
         print(data[start:end])
-        data = data[:start] + '\n' + data[end:]
+        data = data[:start] + b'\n' + data[end:]
 
     return data
 
 
-cautious_check_regex = re.compile(
-    r'\s*import\s*base64;?\s*',
-    re.DOTALL,
-)
+def check_ma_file(file, is_cautious=False):
+    file = os.path.abspath(file).replace('\\', '/')
+    with open(file, 'rb') as f:
+        # data = _delete_ma_byte(f.read())
+        data = f.read()
+
+    if is_cautious and cautious_check_regex.search(data):
+        return True
+    if ma_find_script_node_regex.search(data):
+        return True
+    return False
 
 
-def check_file(file, is_cautious=False):
+def process_ma_file(file):
     file = os.path.abspath(file).replace('\\', '/')
     with open(file, 'rb') as f:
-        data = _delete_ma_byte(f.read())
+        # data = _delete_ma_byte(f.read())
+        data = f.read()
+
+    new_data = _process_ma_content(data)
+    if new_data != data:
+        with open(file, 'wb') as f:
+            f.write(new_data)
+    else:
+        cmds.warning('"{}" 未发现病毒'.format(file))
 
+
+def check_mb_file(file, is_cautious=False):
+    file = os.path.abspath(file).replace('\\', '/')
+    with open(file, 'rb') as f:
+        # data = _delete_ma_byte(f.read())
+        data = f.read()
     if is_cautious and cautious_check_regex.search(data):
         return True
-    if find_script_node_regex.search(data):
+    if mb_regex.search(data):
         return True
     return False
 
 
-def _delete_ma_byte(data):
-    try:
-        data = data.decode('utf-8')
-    except UnicodeDecodeError:
-        try:
-            data = data.decode('gbk')
-        except UnicodeDecodeError:
-            raise ValueError('无法解码文件')
+def _process_mb_content(data):
+    for i in reversed(list(mb_regex.finditer(data))):
+        start, end = i.span()
+        print(data[start:end])
+        data = data[:start] + b'\n' + data[end:]
+
     return data
 
 
-def process_file(file):
+def process_mb_file(file):
     file = os.path.abspath(file).replace('\\', '/')
     with open(file, 'rb') as f:
-        data = _delete_ma_byte(f.read())
+        data = f.read()
 
-    new_data = _process_str(data)
+    new_data = _process_mb_content(data)
     if new_data != data:
         with open(file, 'wb') as f:
             f.write(new_data)
@@ -74,7 +101,18 @@ def process_file(file):
 
 if __name__ == '__main__':
     with open('./../testfiles/key.ma', 'rb') as f:
-        # print(_process_str(f.read()))
-        _test_data = _delete_ma_byte(f.read())
-        assert _process_str(_test_data) != _test_data
-    assert check_file('./../testfiles/key.ma', is_cautious=True)
+        _test_data = f.read()
+        assert _process_ma_content(_test_data) != _test_data
+    with open('./../testfiles/key.mb', 'rb') as f:
+        _test_data = f.read()
+        assert _process_mb_content(_test_data) != _test_data
+    with open('./../testfiles/normal.ma', 'rb') as f:
+        _test_data = f.read()
+        assert _process_ma_content(_test_data) == _test_data
+    with open('./../testfiles/normal.mb', 'rb') as f:
+        _test_data = f.read()
+        assert _process_mb_content(_test_data) == _test_data
+    assert check_ma_file('./../testfiles/key.ma', is_cautious=True)
+    assert check_mb_file('./../testfiles/key.mb', is_cautious=True)
+    assert not check_ma_file('./../testfiles/normal.ma', is_cautious=True)
+    assert not check_mb_file('./../testfiles/normal.mb', is_cautious=True)