diff --git a/libbs/__init__.py b/libbs/__init__.py index 852ca36..2540447 100644 --- a/libbs/__init__.py +++ b/libbs/__init__.py @@ -1,4 +1,4 @@ -__version__ = "3.2.0" +__version__ = "3.3.0" import logging diff --git a/libbs/decompilers/angr/interface.py b/libbs/decompilers/angr/interface.py index 8b8452a..334750b 100644 --- a/libbs/decompilers/angr/interface.py +++ b/libbs/decompilers/angr/interface.py @@ -12,7 +12,8 @@ DecompilerInterface, ) from libbs.artifacts import ( - Function, FunctionHeader, Comment, StackVariable, FunctionArgument, Artifact, Decompilation, Context + Function, FunctionHeader, Comment, StackVariable, FunctionArgument, Artifact, Decompilation, Context, + Struct, StructMember ) from .artifact_lifter import AngrArtifactLifter @@ -256,8 +257,9 @@ def _set_function(self, func: Function, **kwargs) -> bool: # re-decompile a function if needed decompilation = self.decompile_function(angr_func).codegen changes = super()._set_function(func, decompilation=decompilation, **kwargs) - if not self.headless: - self.refresh_decompilation(func.addr) + if not self.headless and changes: + # Use "retype_variable" event to trigger proper UI refresh including type reflow + self.refresh_decompilation(func.addr, event="retype_variable") return changes @@ -348,12 +350,34 @@ def _set_stack_variable(self, svar: StackVariable, decompilation=None, **kwargs) return changed dec_svar = AngrInterface.find_stack_var_in_codegen(decompilation, svar.offset) - if dec_svar and svar.name and svar.name != dec_svar.name: - # TODO: set the types of the stack vars + if not dec_svar: + return changed + + # Set the name if provided and different + if svar.name and svar.name != dec_svar.name: dec_svar.name = svar.name dec_svar.renamed = True changed = True + # Set the type if provided + if svar.type: + try: + from angr.sim_type import parse_type + types_store = self.main_instance.project.kb.types + arch = self.main_instance.project.arch + + # Parse the type string into a SimType + sim_type = parse_type(svar.type, predefined_types=types_store, arch=arch) + sim_type = sim_type.with_arch(arch) + + # Get the variable manager and set the type + variable_kb = decompilation._variable_kb if hasattr(decompilation, '_variable_kb') else self.main_instance.project.kb + variable_manager = variable_kb.variables[svar.addr] + variable_manager.set_variable_type(dec_svar, sim_type, all_unified=True, mark_manual=True) + changed = True + except Exception as e: + l.warning(f"Failed to set stack variable type for {svar.name}: {e}") + return changed def _set_comment(self, comment: Comment, decompilation=None, **kwargs) -> bool: @@ -382,6 +406,109 @@ def _set_comment(self, comment: Comment, decompilation=None, **kwargs) -> bool: func_addr = comment.func_addr or self.get_closest_function(comment.addr) return changed & self.refresh_decompilation(func_addr) + # structs + def _structs(self) -> Dict[str, Struct]: + """ + Returns a dict of libbs.Struct that contain the name and size of each struct in the decompiler. + """ + from angr.sim_type import SimStruct, TypeRef + structs = {} + types_store = self.main_instance.project.kb.types + + for type_ref in types_store.iter_own(): + if not isinstance(type_ref, TypeRef): + continue + sim_type = type_ref.type + if isinstance(sim_type, SimStruct): + structs[type_ref.name] = Struct(type_ref.name, sim_type.size // 8 if sim_type.size else 0, {}) + + return structs + + def _get_struct(self, name) -> Optional[Struct]: + """ + Get a struct by name from the TypesStore. + """ + from angr.sim_type import SimStruct, TypeRef + types_store = self.main_instance.project.kb.types + + try: + type_ref = types_store[name] + except KeyError: + return None + + if not isinstance(type_ref, TypeRef): + return None + + sim_struct = type_ref.type + if not isinstance(sim_struct, SimStruct): + return None + + return self._angr_struct_to_libbs(name, sim_struct) + + def _set_struct(self, struct: Struct, header=True, members=True, **kwargs) -> bool: + """ + Create or update a struct in the TypesStore. + """ + from angr.sim_type import SimStruct, TypeRef, parse_type + from collections import OrderedDict + + types_store = self.main_instance.project.kb.types + arch = self.main_instance.project.arch + + # Build the fields OrderedDict from LibBS struct members + fields = OrderedDict() + if members and struct.members: + sorted_members = sorted(struct.members.items(), key=lambda x: x[0]) + for offset, member in sorted_members: + # Parse the member type string into a SimType + try: + sim_type = parse_type(member.type, predefined_types=types_store, arch=arch) + except Exception: + # Fallback to a simple int type with the right size if parsing fails + from angr.sim_type import SimTypeInt + sim_type = SimTypeInt(signed=False).with_arch(arch) + + fields[member.name] = sim_type.with_arch(arch) + + # Create the SimStruct + sim_struct = SimStruct(fields, name=struct.name, pack=True) + sim_struct = sim_struct.with_arch(arch) + + # Wrap it in a TypeRef and store it + type_ref = TypeRef(struct.name, sim_struct) + types_store[struct.name] = type_ref + + return True + + def _del_struct(self, name) -> bool: + """ + Delete a struct from the TypesStore. + """ + types_store = self.main_instance.project.kb.types + + if name in types_store.data: + del types_store.data[name] + return True + + return False + + @staticmethod + def _angr_struct_to_libbs(name: str, sim_struct: "angr.sim_type.SimStruct") -> Struct: + """ + Convert an angr SimStruct to a LibBS Struct. + """ + members = {} + if sim_struct._arch is not None: + offsets = sim_struct.offsets + for field_name, sim_type in sim_struct.fields.items(): + offset = offsets.get(field_name, 0) + type_str = sim_type.c_repr() if sim_type else None + size = sim_type.size // 8 if sim_type and sim_type.size else 0 + members[offset] = StructMember(field_name, offset, type_str, size) + + size = sim_struct.size // 8 if sim_struct.size else 0 + return Struct(name, size, members) + # # Utils # @@ -431,13 +558,16 @@ def addr_starts_instruction(self, addr) -> bool: return addr in node.instruction_addrs - def refresh_decompilation(self, func_addr): + def refresh_decompilation(self, func_addr, event=None): if self.headless: return False self.workspace.jump_to(func_addr) view = self.workspace._get_or_create_view("pseudocode", CodeView) - view.codegen.am_event() + if event: + view.codegen.am_event(event=event) + else: + view.codegen.am_event() view.focus() return True diff --git a/tests/test_decompilers.py b/tests/test_decompilers.py index 42c4e61..7b7dbbc 100644 --- a/tests/test_decompilers.py +++ b/tests/test_decompilers.py @@ -436,6 +436,7 @@ def test_angr(self): headless=True, binary_path=self.FAUXWARE_PATH ) + self.deci = deci func_addr = deci.art_lifter.lift_addr(0x400664) main = deci.functions[func_addr] main.name = self.RENAMED_NAME @@ -443,6 +444,92 @@ def test_angr(self): assert deci.functions[func_addr].name == self.RENAMED_NAME assert self.RENAMED_NAME in deci.main_instance.project.kb.functions + # + # Struct support + # + + # test struct creation + new_struct = Struct() + new_struct.name = "my_angr_struct" + new_struct.add_struct_member('char_member', 0, 'char', 1) + new_struct.add_struct_member('int_member', 1, 'int', 4) + deci.structs[new_struct.name] = new_struct + + # verify struct was created + updated = deci.structs[new_struct.name] + assert updated is not None, "Struct was not created" + assert updated.name == new_struct.name + + # verify members are present + assert 0 in updated.members, "First member not found" + assert 1 in updated.members, "Second member not found" + + # test struct listing + struct_items = list(deci.structs.items()) + struct_names = [k for k, v in struct_items] + assert new_struct.name in struct_names, "Struct not found in listing" + + # + # Stack variable type setting + # + + # Get the main function which has stack variables + main_func_addr = deci.art_lifter.lift_addr(0x40071d) + main_func = deci.functions[main_func_addr] + + # Check that we have stack variables + assert len(main_func.stack_vars) > 0, "No stack variables found in main function" + + # Get the first stack variable and change its type to a primitive + first_offset = list(main_func.stack_vars.keys())[0] + original_svar = main_func.stack_vars[first_offset] + + # Set a new type (change to int) + original_svar.type = "int" + deci.functions[main_func_addr] = main_func + + # Verify the type was set by re-fetching the function + updated_func = deci.functions[main_func_addr] + updated_svar = updated_func.stack_vars.get(first_offset) + assert updated_svar is not None, "Stack variable not found after update" + # The type should contain "int" (may be formatted differently by angr) + assert "int" in updated_svar.type.lower() if updated_svar.type else False, \ + f"Stack variable type was not updated to int, got: {updated_svar.type}" + + # + # Stack variable type setting with struct type + # + + # Re-fetch the function to get fresh stack variables + main_func = deci.functions[main_func_addr] + + # Get a stack variable (use the same one or another if available) + svar_offsets = list(main_func.stack_vars.keys()) + struct_test_offset = svar_offsets[0] if len(svar_offsets) == 1 else svar_offsets[1] + struct_test_svar = main_func.stack_vars[struct_test_offset] + + # Set the type to a pointer to our struct + struct_ptr_type = f"struct {new_struct.name} *" + struct_test_svar.type = struct_ptr_type + deci.functions[main_func_addr] = main_func + + # Verify the struct type was set + updated_func = deci.functions[main_func_addr] + updated_svar = updated_func.stack_vars.get(struct_test_offset) + assert updated_svar is not None, "Stack variable not found after struct type update" + assert updated_svar.type is not None, "Stack variable type is None after struct type update" + # The type should contain the struct name + assert new_struct.name in updated_svar.type, \ + f"Stack variable type was not updated to struct pointer, got: {updated_svar.type}" + + # Now test struct deletion (after we're done using it for stack var types) + del deci.structs[new_struct.name] + struct_items = list(deci.structs.items()) + struct_keys = [k for k, v in struct_items] + assert new_struct.name not in struct_keys, "Struct was not deleted" + + deci.shutdown() + def test_binja(self): deci = DecompilerInterface.discover( force_decompiler=BINJA_DECOMPILER,