Add support for angr management (#25)

* Add support for angr management

* Fix versioning

* Fix readme

README.md

@@ -110,4 +110,7 @@ You can find a demo of VarBERT running inside DAILA below:
 ![](./assets/binja_daila.png)
 
 - Ghidra
-![](./assets/ghidra_daila.png)
+![](./assets/ghidra_daila.png)
+
+- angr management
+![](./assets/am_daila.png)

dailalib/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.1.4"
+__version__ = "2.2.0"
 
 from .api import AIAPI, OpenAIAPI
 from libbs.api import DecompilerInterface
@@ -23,7 +23,7 @@ def create_plugin(*args, **kwargs):
     # VarModel API (local variable renaming)
     #
 
-    from varbert import VariableRenamingAPI
+    from varbert.api import VariableRenamingAPI
     var_api = VariableRenamingAPI(delay_init=True)
     # add single interface, which is to rename variables
     gui_ctx_menu_actions["DAILA/VarBERT/varbert_rename_vars"] = ("Suggest new variable names", var_api.query_model)

dailalib/api/ai_api.py

@@ -1,5 +1,6 @@
 from typing import Dict, Optional
 from functools import wraps
+import threading
 
 from libbs.api import DecompilerInterface
 

dailalib/binsync_plugin/ai_user_config_ui.py

@@ -165,7 +165,6 @@ def threaded_add_ai_user_to_project(self):
             globals()['workspace'] = self._controller.workspace
             arch = self._controller.main_instance.project.arch.name
 
-        #import remote_pdb; remote_pdb.RemotePdb("localhost", 4444).set_trace()
         add_ai_user_to_project(
             self.api_key, self.binary_path, self.project_path, username=self.username,
             base_on=self.base_on, headless=True if self.decompiler_backend else False, copy_proj=True, model=self.model,

dailalib/daila_plugin.py

@@ -36,9 +36,25 @@ def create_plugin(*args, **kwargs):
         has_ida = True
     except ImportError:
         has_ida = False
+    try:
+        import angrmanagement
+        has_angr = True
+    except ImportError:
+        has_angr = False
 
-    if not has_ida:
+    if not has_ida and not has_angr:
         create_plugin()
+    elif has_angr:
+        from angrmanagement.plugins import BasePlugin
+        class AngrBSPluginThunk(BasePlugin):
+            def __init__(self, workspace):
+                super().__init__(workspace)
+                globals()["workspace"] = workspace
+                self.plugin = create_plugin()
+
+            def teardown(self):
+                pass
+
 
 def PLUGIN_ENTRY(*args, **kwargs):
     """

dailalib/installer.py

@@ -56,6 +56,11 @@ def install_angr(self, path=None, interactive=True):
         if not path:
             return
 
+        path = path / "DAILA"
+        path.mkdir(parents=True, exist_ok=True)
+        src = self.pkg_path / "plugin.toml"
+        dst = Path(path) / "plugin.toml"
+        self.link_or_copy(src, dst, symlink=True)
         self._copy_plugin_to_path(path)
         return path
 

dailalib/plugin.toml

@@ -0,0 +1,13 @@
+[meta]
+plugin_metadata_version = 0
+
+[plugin]
+name = "DAILA_plugin"
+shortname = "DAILA_plugin"
+version = "0.0.0"
+description = ""
+long_description = "The Decompiler Artificial Intelligence Assistant (DAILA)."
+platforms = ["windows", "linux", "macos"]
+min_angr_version = "9.0.0.0"
+author = "mahaloz"
+entrypoints = ["daila_plugin.py"]

setup.cfg

@@ -0,0 +1,28 @@
+[metadata]
+name = dailalib
+version = attr: dailalib.__version__
+url = https://github.com/mahaloz/DAILA
+classifiers =
+    License :: OSI Approved :: BSD License
+    Programming Language :: Python :: 3
+    Programming Language :: Python :: 3.6
+license = BSD 2 Clause
+license_files = LICENSE
+description = The Decompiler Artificial Intelligence Language Assistant (DAILA) is a tool for adding AI to decompilers.
+long_description = file: README.md
+long_description_content_type = text/markdown
+
+[options]
+install_requires =
+    openai>=1.0.0
+    libbs>=0.9.0
+    tiktoken
+    varbert>=2.0.5
+
+python_requires = >= 3.6
+include_package_data = True
+packages = find:
+
+[options.entry_points]
+console_scripts =
+    daila = dailalib.__main__:main

setup.py

@@ -0,0 +1,3 @@
+from setuptools import setup
+
+setup()