diff --git a/rigify/feature_sets.py b/rigify/feature_sets.py index 9ee6821aa376750f9d935e9b2bfc14a89ca8681c..64f40ba9789c561dc011f1bbc339f58d369ec584 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 18d8e5505c2d145b9e83ac778122a356e1cf0c20..2c8adac7aad48ccfa3257c6fa99570af0cefc6c0 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 c1846b998c55dfe541983cfd6703462b094cbbe6..fac88f1a8eaccaafd8a8ff20e406cb6b0edabd65 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)