Fix prompted data getting deleted on re-query in IDA

mahaloz · Oct 14, 2024 · 2ebe2a6 · 2ebe2a6
1 parent 7b15c74
commit 2ebe2a6
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 10 deletions.
diff --git a/dailalib/__init__.py b/dailalib/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "3.10.1"
+__version__ = "3.10.2"
 
 import os
 # stop LiteLLM from querying at all to the remote server

diff --git a/dailalib/api/litellm/prompts/prompt.py b/dailalib/api/litellm/prompts/prompt.py
@@ -155,9 +155,9 @@ def rename_variables(result, function, ai_api: "AIAPI", **kwargs):
         ai_api._dec_interface.rename_local_variables_by_names(function, result)
 
     @staticmethod
-    def comment_function(result, function, ai_api: "AIAPI", **kwargs):
+    def comment_function(result, function: Function, ai_api: "AIAPI", **kwargs):
         curr_cmt_obj = ai_api._dec_interface.comments.get(function.addr, None)
-        curr_cmt = curr_cmt_obj.comment + "\n" if curr_cmt_obj is not None else ""
+        curr_cmt = curr_cmt_obj.comment + "\n\n" if curr_cmt_obj is not None else ""
 
         ai_api._dec_interface.comments[function.addr] = Comment(
             addr=function.addr,
@@ -200,11 +200,7 @@ def comment_vulnerability(result, function, ai_api: "AIAPI", **kwargs):
             new_num = str(_n + bs_cmt_lines - 2)
             rendered = rendered.replace(num, new_num)
 
-        ai_api._dec_interface.comments[function.addr] = Comment(
-            addr=function.addr,
-            comment=rendered,
-            func_addr=function.addr
-        )
+        Prompt.comment_function(rendered, function, ai_api)
 
     @staticmethod
     def comment_man_page(result, function, ai_api: "AIAPI", context=None, **kwargs):
@@ -221,9 +217,11 @@ def comment_man_page(result, function, ai_api: "AIAPI", context=None, **kwargs):
             rendered = str(result)
 
         addr = context.addr if isinstance(context, Context) and context.addr is not None else function.addr
+        curr_cmt_obj = ai_api._dec_interface.comments.get(addr, None)
+        curr_cmt = curr_cmt_obj.comment + "\n" if curr_cmt_obj is not None else ""
         ai_api._dec_interface.comments[addr] = Comment(
             addr=addr,
-            comment=rendered,
+            comment=curr_cmt + rendered,
             func_addr=function.addr,
             decompiled=True
         )
diff --git a/setup.cfg b/setup.cfg
@@ -17,7 +17,7 @@ install_requires =
     litellm>=1.44.27
     tiktoken
     Jinja2
-    libbs>=2.1.0
+    libbs>=2.2.1
 
 python_requires = >= 3.10
 include_package_data = True