diff options
author | Alexander Gavrilov <angavrilov@gmail.com> | 2019-05-03 20:23:05 +0300 |
---|---|---|
committer | Alexander Gavrilov <angavrilov@gmail.com> | 2019-05-03 20:23:05 +0300 |
commit | 010e955654483c4e5de63d399bd15c5ec4def2d2 (patch) | |
tree | 1838e5834287d844e56dfb00322e5ea44623a0aa /rigify | |
parent | 17de4c60631c5c9e68032a6aacd6003f92513765 (diff) |
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.
Diffstat (limited to 'rigify')
-rw-r--r-- | rigify/feature_sets.py | 38 | ||||
-rw-r--r-- | rigify/metarig_menu.py | 24 | ||||
-rw-r--r-- | rigify/rig_lists.py | 23 |
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) |