Fix stack var setting

libbs/api/artifact_lifter.py

@@ -1,7 +1,7 @@
 import logging
 import typing
 
-from libbs.artifacts import StackVariable, Artifact
+from libbs.artifacts import StackVariable, Artifact, FunctionArgument, StructMember
 from libbs.api.type_parser import CTypeParser
 
 if typing.TYPE_CHECKING:
@@ -105,9 +105,10 @@ def _lift_or_lower_artifact(self, artifact, mode):
                 setattr(lifted_art, attr, self._lift_or_lower_artifact(attr_val, mode))
             # nested args, stack_vars, or struct_members
             elif isinstance(attr_val, dict):
-                nested_arts = {
-                    k: self._lift_or_lower_artifact(v, mode) for k, v in attr_val.items()
-                }
+                nested_arts = {}
+                for k, v in attr_val.items():
+                    nested_art = self._lift_or_lower_artifact(v, mode)
+                    nested_arts[nested_art.offset if isinstance(nested_art, (StackVariable, FunctionArgument, StructMember)) else k] = nested_art
                 setattr(lifted_art, attr, nested_arts)
 
         return lifted_art

libbs/decompilers/ida/artifact_lifter.py

@@ -3,32 +3,46 @@
 from libbs.api import ArtifactLifter
 
 l = logging.getLogger(name=__name__)
-from . import compat
-
 
 class IDAArtifactLifter(ArtifactLifter):
     lift_map = {
         "__int64": "long long",
         "__int32": "int",
         "__int16": "short",
         "__int8": "char",
+        "_BOOL": "bool",
+        "_BOOL1": "bool",
+        "_BOOL2": "short",
+        "_BOOL4": "int",
+        "_BOOL8": "long long",
+        "_BYTE": "char",
+        "_WORD": "unsigned short",
+        "_DWORD": "unsigned int",
+        "_QWORD": "unsigned long long",
     }
 
     def __init__(self, deci):
         super(IDAArtifactLifter, self).__init__(deci)
 
     def lift_type(self, type_str: str) -> str:
-        for ida_t, bs_t in self.lift_map.items():
-            type_str = type_str.replace(ida_t, bs_t)
-
-        return type_str
+        return self.lift_ida_type(type_str)
 
     def lift_stack_offset(self, offset: int, func_addr: int) -> int:
+        from . import compat
         return compat.ida_to_bs_stack_offset(func_addr, offset)
 
     def lower_type(self, type_str: str) -> str:
+        # no need to lift type for IDA, since it parses normal C
+        # types by default
         return type_str
 
     def lower_stack_offset(self, offset: int, func_addr: int) -> int:
-        # TODO: this needs to be updated
-        return abs(offset) #compat.ida_to_angr_stack_offset(func_addr, offset)
+        from . import compat
+        return compat.bs_to_ida_stack_offset(self.lower_addr(func_addr), offset)
+
+    @staticmethod
+    def lift_ida_type(type_str: str) -> str:
+        for ida_t, bs_t in IDAArtifactLifter.lift_map.items():
+            type_str = type_str.replace(ida_t, bs_t)
+
+        return type_str

libbs/decompilers/ida/compat.py

@@ -25,6 +25,7 @@
     StructMember
 )
 
+from .artifact_lifter import IDAArtifactLifter
 if typing.TYPE_CHECKING:
     from .interface import IDAInterface
 
@@ -441,10 +442,7 @@ def set_function(func: Function, decompiler_available=True, **kwargs):
 
     # stack vars
     if func.stack_vars and ida_code_view is not None:
-        for svar in func.stack_vars.values():
-            changes |= set_stack_variable(
-                svar, decompiler_available=decompiler_available, ida_code_view=ida_code_view, **kwargs
-            )
+        changes |= set_stack_variables(func.stack_vars, ida_code_view=ida_code_view)
 
     if changes and ida_code_view is not None:
         ida_code_view.refresh_view(changes)
@@ -708,76 +706,109 @@ def get_func_stack_tif(func):
 
     return tif
 
-
-def ida_to_bs_stack_offset(func_addr: int, ida_stack_off: int):
-    if not new_ida_typing_system():
-        return _deprecated_ida_to_bs_offset(func_addr, ida_stack_off)
-
+def get_frame_info(func_addr) -> typing.Tuple[int, int]:
     func = idaapi.get_func(func_addr)
     if not func:
         raise ValueError(f"Function {hex(func_addr)} does not exist.")
 
     stack_tif = get_func_stack_tif(func)
     if stack_tif is None:
         _l.warning(f"Function {hex(func_addr)} does not have a stack frame.")
-        return ida_stack_off
+        return None, None
 
     frame_size = stack_tif.get_size()
     if frame_size == 0:
         _l.warning(f"Function {hex(func_addr)} has a stack frame size of 0.")
-        return ida_stack_off
+        return None, None
 
     # get the last member size
     udt_data = ida_typeinf.udt_type_data_t()
     stack_tif.get_udt_details(udt_data)
     membs = [m for m in udt_data]
     if not membs:
         _l.warning(f"Function {hex(func_addr)} has a stack frame with no members.")
-        return ida_stack_off
+        return None, None
 
     last_member_type = membs[-1].type
     if not last_member_type:
         _l.warning(f"Function {hex(func_addr)} has a stack frame with a member with no type.")
-        return ida_stack_off
+        return None, None
 
     last_member_size = last_member_type.get_size()
+    return frame_size, last_member_size
+
+def ida_to_bs_stack_offset(func_addr: int, ida_stack_off: int):
+    if not new_ida_typing_system():
+        return _deprecated_ida_to_bs_offset(func_addr, ida_stack_off)
+
+    frame_size, last_member_size = get_frame_info(func_addr)
+    if frame_size is None or last_member_size is None:
+        return ida_stack_off
+
     bs_soff = ida_stack_off - frame_size + last_member_size
     return bs_soff
 
+def bs_to_ida_stack_offset(func_addr: int, bs_stack_off: int):
+    if not new_ida_typing_system():
+        # maintain backwards compatibility
+        return bs_stack_off
 
-@execute_write
-@requires_decompilation
-def set_stack_variable(svar: StackVariable, decompiler_available=True, **kwargs):
+    frame_size, last_member_size = get_frame_info(func_addr)
+    if frame_size is None or last_member_size is None:
+        return ida_stack_off
+
+    ida_soff = bs_stack_off + frame_size - last_member_size
+    return ida_soff
+
+def set_stack_variables(svars: list[StackVariable], decompiler_available=True, **kwargs) -> bool:
+    """
+    This function should only be called in a function that is already used in main-thread.
+    This should also mean decompilation is passed in.
+    """
     ida_code_view = kwargs.get('ida_code_view', None)
     changes = False
-    # if we have a decompiler, we can set the type and name ind decompilation
-    if ida_code_view is not None:
-        if svar.type:
-            ida_type = convert_type_str_to_ida_type(svar.type)
-            if ida_type is None:
-                _l.warning(f"IDA Failed to parse type for stack var {svar}")
-                return changes
-
-            changes |= set_stack_vars_types({svar.offset: ida_type}, ida_code_view)
-            if changes:
-                ida_code_view.cfunc.refresh_func_ctext()
+    if ida_code_view is None:
+        # TODO: support decompilation-less stack var setting
+        _l.warning("Cannot set stack variables without a decompiler.")
+        return changes
+
+    lvars = {v.location.stkoff(): v for v in ida_code_view.cfunc.lvars if v.is_stk_var()}
+    if not lvars:
+        _l.warning("No stack variables found in decompilation to set. Making new ones is not supported")
+        return changes
+
+    for bs_off, bs_var in svars.items():
+        if bs_off not in lvars:
+            _l.warning(f"Stack variable at offset {bs_off} not found in decompilation.")
+            continue
 
-        lvars = [v for v in ida_code_view.cfunc.lvars if v.is_stk_var()]
-        if not lvars:
-            return changes
+        lvar = lvars[bs_off]
 
-        if svar.name:
-            for lvar in lvars:
-                if lvar.location.stkoff() == svar.offset:
-                    lvar.name = svar.name
+        # naming:
+        if bs_var.name and bs_var.name != lvar.name:
+            ida_code_view.rename_lvar(lvar, bs_var.name, 1)
+            changes |= True
+            ida_code_view.cfunc.refresh_func_ctext()
+            lvars = {v.location.stkoff(): v for v in ida_code_view.cfunc.lvars if v.is_stk_var()}
+
+        # typing
+        if bs_var.type:
+            curr_ida_type_str = str(lvar.type()) if lvar.type() else None
+            curr_ida_type = IDAArtifactLifter.lift_ida_type(curr_ida_type_str) if curr_ida_type_str else None
+            if curr_ida_type and bs_var.type != curr_ida_type:
+                new_type = convert_type_str_to_ida_type(bs_var.type)
+                if new_type is None:
+                    _l.warning(f"Failed to convert type string {bs_var.type} to ida type.")
+                    continue
+
+                updated_type = ida_code_view.set_lvar_type(lvar, new_type)
+                if updated_type:
                     changes |= True
-                    break
+                    ida_code_view.cfunc.refresh_func_ctext()
+                    lvars = {v.location.stkoff(): v for v in ida_code_view.cfunc.lvars if v.is_stk_var()}
 
-            if changes:
-                ida_code_view.cfunc.refresh_func_ctext()
-    else:
-        # TODO: support stack vars without a decompiler available
-        _l.warning("Setting stack variables without a decompiler is not supported yet.")
+    if changes:
+        ida_code_view.refresh_view(True)
 
     return changes
 
@@ -814,7 +845,7 @@ def get_func_stack_var_info(func_addr) -> typing.Dict[int, StackVariable]:
 
 
 @execute_write
-def set_stack_vars_types(var_type_dict, ida_code_view) -> bool:
+def _deprecated_set_stack_vars_types(var_type_dict, ida_code_view) -> bool:
     """
     Sets the type of a stack variable, which should be a local variable.
     Take special note of the types of first two parameters used here:

libbs/decompilers/ida/interface.py

@@ -287,7 +287,8 @@ def _functions(self) -> Dict[int, Function]:
 
     # stack vars
     def _set_stack_variable(self, svar: StackVariable, **kwargs) -> bool:
-        return compat.set_stack_variable(svar, headless=self.headless, decompiler_available=self.decompiler_available, **kwargs)
+        _l.warning("Setting stack vars using this API is deprecared. Use _set_function instead.")
+        return False
 
     # global variables
     def _set_global_variable(self, gvar: GlobalVariable, **kwargs) -> bool: