avoid re-importing

main.py

@@ -3,27 +3,41 @@ import subprocess
 import tempfile
 import os
 import sys
+import hashlib
 import json
 
+GHIDRA_PROJECT_DIR = f"{os.getenv('HOME')}/ghidra_project"
+
+os.makedirs(GHIDRA_PROJECT_DIR, exist_ok=True)
+
+def hash_file(file):
+    sha256_hash = hashlib.sha256()
+    with open(file, "rb") as f:
+        for byte_block in iter(lambda: f.read(4096), b""):
+            sha256_hash.update(byte_block)
+    return sha256_hash.hexdigest()
 
 def get_functions(file):
 
+    file_hash = hash_file(file)
+
     with tempfile.TemporaryDirectory() as TEMP_DIR:
 
-        try:
-            output = subprocess.check_output(
-            f"/ghidra/support/analyzeHeadless {TEMP_DIR} Project -import {file} -postscript /home/user/app/scripts/dump_functions.py {TEMP_DIR}/funcs.json",
+        o = subprocess.run(
+            f"/ghidra/support/analyzeHeadless {GHIDRA_PROJECT_DIR} {file_hash} -import {file} -postscript /home/user/app/scripts/dump_functions.py {TEMP_DIR}/funcs.json 2>&1",
             shell=True,
+            capture_output=True,
+            encoding="utf8"
             )
-        except Exception as e:
-            raise gr.Error(f"Unable to run Ghidra on {file}: {e}")
+        if o.returncode != 0:
+            raise gr.Error(f"Unable to run Ghidra on {file}: {o.stdout}")
 
         if not os.path.exists(f"{TEMP_DIR}/funcs.json"):
-                raise gr.Error(f"DIRTY Ghidra failed to produce output: {output}")
+                raise gr.Error(f"DIRTY Ghidra failed to produce output: {o.stdout}")
 
         json_funcs = json.load(open(f"{TEMP_DIR}/funcs.json"))
 
-        return json_funcs
+    return json_funcs
 
 
 with gr.Blocks() as demo:
@@ -38,7 +52,6 @@ with gr.Blocks() as demo:
     To get started, upload a binary or select one of the example binaries below.
 
     ## TODOs
-    * Avoid re-running Ghidra when changing the function
     * Make predictions for variables in non-unique storage locations
     """
     )
@@ -67,19 +80,23 @@ with gr.Blocks() as demo:
 
         with gr.Row(visible=True) as result:
             disassembly = gr.Code(
-                label="Disassembly", lines=20
+                label="Disassembly", lines=20,
+                #min_width=400
             )
             original_decompile = gr.Code(
                 language="c",
-                label="Original Decompilation", lines=20
+                label="Original Decompilation", lines=20,
+                #min_width=400
             )
             decompile = gr.Code(
                 language="c",
                 label="Renamed and retyped Decompilation",
                 lines=20,
+                #min_width=400
             )
             model_output = gr.JSON(
-                label="Model Output"
+                label="Model Output",
+                #min_width=400
             )
             # with gr.Column():
             #     clazz = gr.Label()
@@ -119,7 +136,9 @@ with gr.Blocks() as demo:
             return {
                 col: gr.Column(visible=True),
                 fun_dropdown: gr.Dropdown(choices=addrs, value=addrs[0][1]),
-                state: {"file": file, "cfs": cfs},
+                state: {"file": file,
+                        "file_hash": hash_file(file.name),
+                        "cfs": cfs},
             }
 
     def function_change_fn(selected_fun, state, progress=gr.Progress()):
@@ -132,16 +151,17 @@ with gr.Blocks() as demo:
 
             progress(0, desc=f"Running DIRTY Ghidra on {hex(selected_fun)}...")
 
-            try:
-                output = subprocess.check_output(
-                    f"/ghidra/support/analyzeHeadless {TEMP_DIR} Project -import {state['file'].name} -postscript /DIRTY/scripts/DIRTY_infer.py {TEMP_DIR}/funcs.json {selected_fun}",
+            o = subprocess.run(
+                    f"/ghidra/support/analyzeHeadless {GHIDRA_PROJECT_DIR} {state['file_hash']} -process -postscript /DIRTY/scripts/DIRTY_infer.py {TEMP_DIR}/funcs.json {selected_fun} 2>&1",
                     shell=True,
+                    capture_output=True,
+                    encoding="utf8"
                 )
-            except Exception as e:
-                raise gr.Error(f"Unable to run Ghidra: {e}\n{output}")
+            if o.returncode != 0:
+                raise gr.Error(f"Unable to run Ghidra: {o.stdout}")
 
             if not os.path.exists(f"{TEMP_DIR}/funcs.json"):
-                raise gr.Error(f"DIRTY Ghidra failed to produce output: {output}")
+                raise gr.Error(f"DIRTY Ghidra failed to produce output: {o.stdout}")
 
             json_info = json.load(open(f"{TEMP_DIR}/funcs.json"))