diff --git a/varbert/__init__.py b/varbert/__init__.py index f78d484..e0d4da6 100644 --- a/varbert/__init__.py +++ b/varbert/__init__.py @@ -1,4 +1,4 @@ -__version__ = "2.2.1" +__version__ = "2.3.0" import importlib.resources import tarfile diff --git a/varbert/api.py b/varbert/api.py index 8bc58f4..aba4a3c 100644 --- a/varbert/api.py +++ b/varbert/api.py @@ -47,11 +47,8 @@ def predict_variable_names( if function is None and decompilation_text is None: raise ValueError("Must provide either a Function or decompilation text.") if function: - if not function.args and not function.stack_vars: - self.debug(f"{function} has no arguments or stack variables to predict names for.") - return {}, "" if function.size < self._min_func_size: - self.debug(f"{function} is smaller than min size of {self._min_func_size} bytes.") + self.error(f"{function} is smaller than min size of {self._min_func_size} bytes.") return {}, "" # can be None because of the delay init if self._model_interface is None: @@ -70,6 +67,9 @@ def predict_variable_names( preprocessor = DecompilationTextProcessor( decompilation_text, func=function, decompiler=self._dec_interface if use_decompiler else None ) + if preprocessor.failed: + return {}, "" + processed_code, func_args = preprocessor.processed_code, preprocessor.func_args scores, score_origins = self._model_interface.process(processed_code) if scores is None: diff --git a/varbert/text_processor.py b/varbert/text_processor.py index b9f0f92..b6c0ce2 100644 --- a/varbert/text_processor.py +++ b/varbert/text_processor.py @@ -19,6 +19,8 @@ def __init__(self, raw_code, func: Optional[Function] = None, decompiler: Option self._decompiler = decompiler self._func = func + self.failed = False + # updated in process_code self.processed_code = raw_code # TODO: this actually needs to be disabled for now because Function does not handle register variables @@ -83,6 +85,11 @@ def _process_code_with_decompiler(self): # refresh the decompiled obj backend self._func.dec_obj = self._decompiler.get_decompilation_object(self._func) original_names = self._decompiler.local_variable_names(self._func) + if not original_names: + self._decompiler.error(f"No variable names (local or args) found in decompiled function {self._func} for renaming.") + self.failed = True + return + og_name_to_tokenized_name = self._tokenize_names(original_names, token=tmp_token) tokenized_name_to_og_name = {v: k for k, v in og_name_to_tokenized_name.items()}