Use Ghidra to enumerate functions

Dockerfile

@@ -28,8 +28,8 @@ RUN unzip ghidrathon/Ghidrathon-v4.0.0.zip -d /ghidra/Ghidra/Extensions
 
 WORKDIR /
 
-#RUN git clone -b main https://github.com/edmcman/DIRTY
-ADD ./DIRTY /DIRTY
+RUN git clone -b main https://github.com/edmcman/DIRTY
+#ADD ./DIRTY /DIRTY
 
 RUN --mount=type=cache,target=/root/.cache pip install --no-cache-dir --upgrade -r /DIRTY/requirements.txt
 

main.py

@@ -4,22 +4,21 @@ import subprocess
 import tempfile
 import os
 import sys
+import json
 
-def new_binary(file):
+def get_functions(file):
 
-    with tempfile.TemporaryDirectory() as TEMP_DIR, tempfile.TemporaryDirectory() as OUTPUT_DIR:
-        shutil.copy2(file.name, TEMP_DIR)
-        subprocess.run(f"ls -lR {TEMP_DIR}", shell=True)
+    with tempfile.TemporaryDirectory() as TEMP_DIR:
 
-        # python3 generate.py --ghidra PATH_TO_GHIDRA_ANALYZEHEADLESS -t NUM_THREADS -n [NUM_FILES|None] -b BINARIES_DIR -o OUTPUT_DIR
-        print("Running DIRTY-Ghidra...", file=sys.stderr)
-        subprocess.run(f"python /DIRTY/dataset-gen-ghidra/generate.py --verbose --ghidra /ghidra/support/analyzeHeadless -t 1 -b {TEMP_DIR} -o {OUTPUT_DIR}", shell=True)
-        subprocess.run(f"ls -lR {OUTPUT_DIR}", shell=True)
+        subprocess.run(f"/ghidra/support/analyzeHeadless {TEMP_DIR} Project -import {file} -postscript /home/user/app/scripts/dump_functions.py {TEMP_DIR}/funcs.json", shell=True)
 
+        json_funcs = json.load(open(f"{TEMP_DIR}/funcs.json"))
+
+        return json_funcs
 
 with gr.Blocks() as demo:
 
-    #all_dis_state = gr.State()
+    all_dis_state = gr.State()
 
     gr.Markdown(
         """
@@ -30,30 +29,55 @@ with gr.Blocks() as demo:
 
     file_widget = gr.File(label="Executable file")
 
-    def file_change_fn(file, progress=gr.Progress()):
+    with gr.Column(visible=False) as col:
+        #output = gr.Textbox("Output")
+
+        gr.Markdown("""
+        Great, you selected an executable!  Now pick the function you would like to analyze.                    
+                    """)
+
+        fun_dropdown = gr.Dropdown(label="Select a function", choices=["Woohoo!"], interactive=True)
 
+        gr.Markdown("""
+        Below you can find the selected function's disassembly, and the model's
+        prediction of whether the function is an object-oriented method or a
+        regular function.
+                    """)
+
+        with gr.Row(visible=True) as result:
+            disassembly = gr.Textbox(label="Disassembly", lines=20)
+            with gr.Column():
+                clazz = gr.Label()
+                #interpret_button = gr.Button("Interpret (very slow)")
+            #interpretation = gr.components.Interpretation(disassembly)
+
+    def file_change_fn(file, progress=gr.Progress()):
 
         if file is None:
             return {
-                    #col: gr.update(visible=False),
-                    #all_dis_state: None
+                    col: gr.update(visible=False),
+                    all_dis_state: None
                     }
         else:
 
             #fun_data = {42: 2, 43: 3}
-            new_binary(file)
-            progress(0, desc="Decompiling binary...")
-            #fun_data = get_all_dis(file.name)
+            #new_binary(file)
+            progress(0, desc="Analyzing binary...")
+            try:
+                fun_data = get_functions(file.name)
+                #print(fun_data)
 
-            #addrs = ["%#x" % addr for addr in fun_data.keys()]
+                addrs = [(f"{name} ({hex(int(addr))})", int(addr)) for addr, name in fun_data.items()]
+            except:
+                raise gr.Error("Unable to obtain functions")
 
             return {
-                    #col: gr.update(visible=True),
-                    #fun_dropdown: gr.Dropdown.update(choices=addrs, value=addrs[0]),
-                    #all_dis_state: fun_data
+                    col: gr.Column(visible=True),
+                    fun_dropdown: gr.Dropdown(choices=addrs, value=addrs[0][1]),
+                    all_dis_state: fun_data
                     }
         
-    file_widget.change(file_change_fn, file_widget)
+    file_widget.change(file_change_fn, file_widget, outputs=[col, fun_dropdown, all_dis_state])
 
 # spaces only shows stderr..
 os.dup2(sys.stdout.fileno(), sys.stderr.fileno())

scripts/dump_functions.py

@@ -8,8 +8,12 @@ def dump_functions_to_json():
     functions_dict = {}
     
     for func in functions:
+
+        if func.isExternal() or func.isThunk():
+            continue
+
         func_name = func.getName()
-        func_address = func.getEntryPoint().toString()
+        func_address = func.getEntryPoint().getOffset()
         
         # Add function name and address to the dictionary
         functions_dict[func_address] = func_name