From 7e2ac5f78badb957a9c14874538888cbb9438dea Mon Sep 17 00:00:00 2001
From: Alessandro Gario <5714290+alessandrogario@users.noreply.github.com>
Date: Sat, 13 Feb 2021 00:20:47 +0100
Subject: [PATCH] IDA: Add a simple action to generate spec files (#94)

* IDA: Add a simple action to generate spec files

* docs: Update the example instructions
---
 README.md                     | 22 ++++++++--
 examples/ida_plugin_action.py | 78 +++++++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+), 4 deletions(-)
 create mode 100644 examples/ida_plugin_action.py

diff --git a/README.md b/README.md
index 48e71abc6..2192af4c7 100644
--- a/README.md
+++ b/README.md
@@ -85,11 +85,25 @@ These depend on tools like [IDA Pro](https://www.hex-rays.com/products/ida) or [
 
 Given that we have either of the above, we can try out Anvill's machine code lifter on a binary of our choice.
 
+**First, we generate a JSON specification from a binary:**
+
+From the CLI:
+
 ```shell
-# First, we generate a JSON specification from a binary
-python3.8 -m anvill --bin_in my_binary --spec_out spec.json
-# Finally we produce LLVM bitcode from a JSON specification
-./remill-build/tools/anvill/anvill-lift-json-*.0 --spec spec.json --bc_out out.bc
+python3 -m anvill --bin_in my_binary --spec_out spec.json
+```
+
+With the IDA plugin:
+1. Open the binary inside IDA
+2. Select **Run script** in the **File** menu
+3. Open the `examples/ida_plugin_action.py`
+4. In the disasm window, place the cursor inside a function
+5. Right click and select **Generate ANVILL spec file**
+
+**Finally we produce LLVM bitcode from a JSON specification**
+
+```
+./build/anvill-decompile-json-*.0 --spec spec.json --bc_out out.bc
 ```
 
 ### Docker image
diff --git a/examples/ida_plugin_action.py b/examples/ida_plugin_action.py
new file mode 100644
index 000000000..b1c2ce895
--- /dev/null
+++ b/examples/ida_plugin_action.py
@@ -0,0 +1,78 @@
+# Copyright (c) 2021-present Trail of Bits, Inc.
+
+import ida_funcs
+import ida_kernwin
+import idautils
+
+import anvill
+import json
+
+class generate_anvill_spec_t(ida_kernwin.action_handler_t):
+  def activate(self, ctx):
+    user_input = ida_kernwin.ask_yn(ida_kernwin.ASKBTN_YES, "Would you like to export all functions?")
+    if user_input == ida_kernwin.ASKBTN_CANCEL:
+      return 1
+
+    output_file_name_hint = ""
+
+    p = anvill.get_program()
+
+    if user_input == ida_kernwin.ASKBTN_NO:
+      screen_cursor = ida_kernwin.get_screen_ea()
+      function_name = ida_funcs.get_func_name(screen_cursor)
+      if function_name is None:
+        print("ANVILL: The cursor is not located inside a function")
+        return 1
+
+      output_file_name_hint = function_name + ".json"
+
+      try:
+        p.add_function_definition(screen_cursor)
+
+      except:
+        print("ANVILL: Failed to process the function at address {0:x}".format(screen_cursor))
+        return 1
+
+    else:
+      function_address_list = idautils.Functions()
+      for function_address in function_address_list:
+        try:
+          p.add_function_definition(function_address)
+
+        except:
+          print("ANVILL: Failed to process the function at address {0:x}".format(function_address))
+
+      output_file_name_hint = "program.json"
+
+    output_path = ida_kernwin.ask_file(True, output_file_name_hint, "Select where to save the spec file")
+    if not output_path:
+      return 1
+
+    output = json.dumps(p.proto(), sort_keys=False, indent=2)
+
+    print("ANVILL: Saving the spec file to {}".format(output_path))
+    with open(output_path, "w") as f:
+      f.write(output)
+
+  def update(self, ctx):
+    if ctx.widget_type == ida_kernwin.BWN_DISASM:
+      return ida_kernwin.AST_ENABLE_FOR_WIDGET
+
+    return ida_kernwin.AST_DISABLE_FOR_WIDGET
+
+ACTION_NAME = "generate-anvill-spec-file"
+
+ida_kernwin.register_action(
+    ida_kernwin.action_desc_t(
+        ACTION_NAME,
+        "Generate ANVILL spec file",
+        generate_anvill_spec_t(),
+        "Ctrl+H"))
+
+class popup_hooks_t(ida_kernwin.UI_Hooks):
+  def finish_populating_widget_popup(self, w, popup):
+    if ida_kernwin.get_widget_type(w) == ida_kernwin.BWN_DISASM:
+      ida_kernwin.attach_action_to_popup(w, popup, ACTION_NAME, None)
+
+hooks = popup_hooks_t()
+hooks.hook()