Skip to content

Commit

Permalink
Allow renaming of functions without stack or parameters (#13)
Browse files Browse the repository at this point in the history
* Allow renaming of functions without stack or parameters
* return on failure
  • Loading branch information
mahaloz authored Sep 9, 2024
1 parent acb1b42 commit 1e6a85b
Show file tree
Hide file tree
Showing 3 changed files with 12 additions and 5 deletions.
2 changes: 1 addition & 1 deletion varbert/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
__version__ = "2.2.1"
__version__ = "2.3.0"

import importlib.resources
import tarfile
Expand Down
8 changes: 4 additions & 4 deletions varbert/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down
7 changes: 7 additions & 0 deletions varbert/text_processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()}

Expand Down

0 comments on commit 1e6a85b

Please sign in to comment.