Rigify: improve robustness with bad feature set packages.

Verify the basic expected directory structure inside the ZIP
archive before installing it, and catch exceptions when loading
the already installed packages.

3 files changed, 78 insertions, 7 deletions

diff --git a/rigify/feature_sets.py b/rigify/feature_sets.py
index 9ee6821a..64f40ba9 100644
--- a/rigify/feature_sets.py
+++ b/rigify/feature_sets.py
@@ -19,6 +19,7 @@
 import bpy
 from bpy.props import StringProperty
 import os
+import re
 from zipfile import ZipFile
 from shutil import rmtree
 
@@ -34,6 +35,29 @@ def feature_set_items(scene, context):
 
     return items
 
+def verify_feature_set_archive(zipfile):
+    """Verify that the zip file contains one root directory, and some required files."""
+    dirname = None
+    init_found = False
+    data_found = False
+
+    for name in zipfile.namelist():
+        parts = re.split(r'[/\\]', name)
+
+        if dirname is None:
+            dirname = parts[0]
+        elif dirname != parts[0]:
+            dirname = None
+            break
+
+        if len(parts) == 2 and parts[1] == '__init__.py':
+            init_found = True
+
+        if len(parts) > 2 and parts[1] in {'rigs', 'metarigs'} and parts[-1] == '__init__.py':
+            data_found = True
+
+    return dirname, init_found, data_found
+
 class DATA_OT_rigify_add_feature_set(bpy.types.Operator):
     bl_idname = "wm.rigify_add_feature_set"
     bl_label = "Add External Feature Set"
@@ -57,6 +81,20 @@ class DATA_OT_rigify_add_feature_set(bpy.types.Operator):
         rigify_config_path = os.path.join(bpy.utils.script_path_user(), 'rigify')
         os.makedirs(rigify_config_path, exist_ok=True)
         with ZipFile(bpy.path.abspath(self.filepath), 'r') as zip_archive:
+            base_dirname, init_found, data_found = verify_feature_set_archive(zip_archive)
+
+            if not base_dirname:
+                self.report({'ERROR'}, "The feature set archive must contain one base directory.")
+                return {'CANCELLED'}
+
+            if not re.fullmatch(r'[a-zA-Z_][a-zA-Z_0-9-]*', base_dirname):
+                self.report({'ERROR'}, "The feature set archive has invalid characters in the base directory name: '%s'." % (base_dirname))
+                return {'CANCELLED'}
+
+            if not init_found or not data_found:
+                self.report({'ERROR'}, "The feature set archive has no rigs or metarigs, or is missing __init__.py.")
+                return {'CANCELLED'}
+
             zip_archive.extractall(rigify_config_path)
 
         addon_prefs.machin = bpy.props.EnumProperty(items=(('a',)*3, ('b',)*3, ('c',)*3),)
diff --git a/rigify/metarig_menu.py b/rigify/metarig_menu.py
index 18d8e550..2c8adac7 100644
--- a/rigify/metarig_menu.py
+++ b/rigify/metarig_menu.py
@@ -19,6 +19,8 @@
 # <pep8 compliant>
 
 import os
+import traceback
+
 from string import capwords
 
 import bpy
@@ -44,7 +46,11 @@ def get_metarigs(base_path, path, depth=0):
 
     metarigs = {}
 
-    files = os.listdir(os.path.join(base_path, path))
+    try:
+        files = os.listdir(os.path.join(base_path, path))
+    except FileNotFoundError:
+        files = []
+
     files.sort()
 
     for f in files:
@@ -216,9 +222,19 @@ def get_external_metarigs(feature_sets_path):
 
     for feature_set in os.listdir(feature_sets_path):
         if feature_set:
-            utils.get_resource(os.path.join(feature_set, '__init__'), base_path=feature_sets_path)
-
-            metarigs['external'].update(get_metarigs(feature_sets_path, os.path.join(feature_set, utils.METARIG_DIR)))
+            try:
+                try:
+                    utils.get_resource(os.path.join(feature_set, '__init__'), feature_sets_path)
+                except FileNotFoundError:
+                    print("Rigify Error: Could not load feature set '%s': __init__.py not found.\n" % (feature_set))
+                    continue
+
+                metarigs['external'].update(get_metarigs(feature_sets_path, os.path.join(feature_set, utils.METARIG_DIR)))
+            except Exception:
+                print("Rigify Error: Could not load feature set '%s' metarigs: exception occurred.\n" % (feature_set))
+                traceback.print_exc()
+                print("")
+                continue
 
     metarig_ops.clear()
     armature_submenus.clear()
diff --git a/rigify/rig_lists.py b/rigify/rig_lists.py
index c1846b99..fac88f1a 100644
--- a/rigify/rig_lists.py
+++ b/rigify/rig_lists.py
@@ -17,6 +17,7 @@
 #======================= END GPL LICENSE BLOCK ========================
 
 import os
+import traceback
 
 from . import utils
 
@@ -33,7 +34,11 @@ def get_rigs(base_path, path, feature_set='rigify'):
     rigs = {}
     impl_rigs = {}
 
-    files = os.listdir(os.path.join(base_path, path))
+    try:
+        files = os.listdir(os.path.join(base_path, path))
+    except FileNotFoundError:
+        files = []
+
     files.sort()
 
     for f in files:
@@ -84,7 +89,19 @@ def get_external_rigs(feature_sets_path):
     # Get external rigs
     for feature_set in os.listdir(feature_sets_path):
         if feature_set:
-            utils.get_resource(os.path.join(feature_set, '__init__'), feature_sets_path)
-            external_rigs, external_impl_rigs = get_rigs(feature_sets_path, os.path.join(feature_set, utils.RIG_DIR), feature_set)
+            try:
+                try:
+                    utils.get_resource(os.path.join(feature_set, '__init__'), feature_sets_path)
+                except FileNotFoundError:
+                    print("Rigify Error: Could not load feature set '%s': __init__.py not found.\n" % (feature_set))
+                    continue
+
+                external_rigs, external_impl_rigs = get_rigs(feature_sets_path, os.path.join(feature_set, utils.RIG_DIR), feature_set)
+            except Exception:
+                print("Rigify Error: Could not load feature set '%s' rigs: exception occurred.\n" % (feature_set))
+                traceback.print_exc()
+                print("")
+                continue
+
             rigs.update(external_rigs)
             implementation_rigs.update(external_impl_rigs)