Skip to content

Commit

Permalink
Fix new LibBS decompiler interface (#12)
Browse files Browse the repository at this point in the history
* Fix new decomp interface

* bump

* bump libbs version for copy fix
  • Loading branch information
mahaloz authored Aug 7, 2024
1 parent 62c543c commit acb1b42
Show file tree
Hide file tree
Showing 4 changed files with 14 additions and 6 deletions.
2 changes: 1 addition & 1 deletion setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ install_requires =
transformers
tqdm
dailalib
libbs
libbs>=1.18.1

python_requires = >= 3.8
packages = find:
Expand Down
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.0"
__version__ = "2.2.1"

import importlib.resources
import tarfile
Expand Down
9 changes: 6 additions & 3 deletions varbert/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,8 @@ def __init__(self, **kwargs):
) if not self._delay_init else None

def predict_variable_names(
self, function: Function = None, decompilation_text: Optional[str] = None, use_decompiler=True,
remove_bad_names=True
self, function: Function = None, decompilation_text: Optional[str] = None, use_decompiler=True,
remove_bad_names=True, **kwargs
) -> Tuple[Dict[str, str], str]:
"""
Predict variable names for a function or decompilation text. You can use this function in two ways:
Expand Down Expand Up @@ -61,7 +61,10 @@ def predict_variable_names(
if not use_decompiler:
raise ValueError("Must provide decompilation text if not using decompiler")

decompilation_text = self._dec_interface.decompile(function.addr)
decomp = self._dec_interface.decompile(function.addr)
if decomp is None or not decomp.text:
raise ValueError(f"Unable to decompile function {function}")
decompilation_text = decomp.text

# preprocess text for training
preprocessor = DecompilationTextProcessor(
Expand Down
7 changes: 6 additions & 1 deletion varbert/text_processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,12 @@ def _process_code_with_decompiler(self):
# replace all original names with tmp tokenized names
self._decompiler.rename_local_variables_by_names(self._func, og_name_to_tokenized_name)
# get the decomp, fix the tmp tokens
tokenized_dec_text = self._decompiler.decompile(self._func.addr)
tokenized_decomp = self._decompiler.decompile(self._func.addr)
if tokenized_decomp is None:
_l.error("Decompiler failed to decompile function.")
return

tokenized_dec_text = tokenized_decomp.text
tokenized_dec_text = tokenized_dec_text.replace(tmp_token, "@@")

# revert to the original names in the decomp
Expand Down

0 comments on commit acb1b42

Please sign in to comment.