updates

Dockerfile

@@ -11,8 +11,7 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
   apt-get -y update && apt-get -y install -y python3-pip python-is-python3 \
   git pkg-config libsentencepiece-dev nano sudo unzip
 
-# RUN --mount=type=cache,target=/root/.npm
-RUN --mount=type=cache,target=/root/.cache pip install --no-cache-dir --upgrade -r /code/requirements.txt
+RUN --mount=type=cache,target=/root/.cache pip install --upgrade -r /code/requirements.txt
 
 # Install Ghidrathon
 
@@ -20,7 +19,7 @@ WORKDIR /tmp/
 
 RUN wget https://github.com/mandiant/Ghidrathon/releases/download/v4.0.0/Ghidrathon-v4.0.0.zip
 RUN unzip Ghidrathon-v4.0.0.zip -d ghidrathon
-RUN python -m pip install -r ghidrathon/requirements.txt
+RUN --mount=type=cache,target=/root/.cache python -m pip install -r ghidrathon/requirements.txt
 RUN python ghidrathon/ghidrathon_configure.py /ghidra
 RUN unzip ghidrathon/Ghidrathon-v4.0.0.zip -d /ghidra/Ghidra/Extensions
 
@@ -28,10 +27,12 @@ RUN unzip ghidrathon/Ghidrathon-v4.0.0.zip -d /ghidra/Ghidra/Extensions
 
 WORKDIR /
 
-RUN git clone -b main https://github.com/edmcman/DIRTY
+RUN git clone -b main https://github.com/edmcman/DIRTY # 09-20-24
 #ADD ./DIRTY /DIRTY
 
-RUN --mount=type=cache,target=/root/.cache pip install --no-cache-dir --upgrade -r /DIRTY/requirements.txt
+RUN --mount=type=cache,target=/root/.cache pip install --upgrade -r /DIRTY/requirements.txt
+
+RUN --mount=type=cache,target=/root/.cache (cd /root/.cache; (echo "b1e812b758eccf402271607c40fa491b5486742abf3706be174dc3f4fe87b9dd data1.tar.bz2" | sha256sum -c || wget -O '/root/.cache/data1.tar.bz2' 'https://cmu.box.com/shared/static/nx9fyn8jx0i9p4bftw8f2giqlufnoyj5')) && tar -xvjf /root/.cache/data1.tar.bz2 -C /DIRTY/dirty
 
 # Set up a new user named "user" with user ID 1000
 RUN useradd -m -u 1000 user

main.py

@@ -6,78 +6,150 @@ import os
 import sys
 import json
 
+
 def get_functions(file):
 
     with tempfile.TemporaryDirectory() as TEMP_DIR:
 
-        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)
+        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()
+    state = gr.State()
 
     intro = gr.Markdown(
         """
-    # DIRTY-Ghidra Testing
-    First, upload a binary.
+    # DIRTY-Ghidra Inference Demo
+    Welcome!  This is a demo of DIRTY-Ghidra, a tool that predict names and types for variables for Ghidra's decompiler.
+
+    To get started, upload a binary.
     """
     )
 
     file_widget = gr.File(label="Executable file")
 
     with gr.Column(visible=False) as col:
-        #output = gr.Textbox("Output")
+        # output = gr.Textbox("Output")
 
-        gr.Markdown("""
-        Great, you selected an executable!  Now pick the function you would like to analyze.                    
-                    """)
+        gr.Markdown(
+            """
+        Great, you selected an executable!  Now pick the function you would like
+        to analyze.  Simple functions (without variables) will probably fail, so
+        you may have to try a few before you find one that works.
+        """
+        )
 
-        fun_dropdown = gr.Dropdown(label="Select a function", choices=["Woohoo!"], interactive=True)
+        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.
-                    """)
+        gr.Markdown(
+            """
+        Below you can find some information.
+                    """
+        )
 
         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)
+            disassembly = gr.Textbox(
+                label="Disassembly", value="Please wait...", lines=20
+            )
+            original_decompile = gr.Textbox(
+                label="Original Decompilation", value="Please wait...", lines=20
+            )
+            decompile = gr.Textbox(
+                label="Renamed and retyped Decompilation",
+                value="Please wait...",
+                lines=20,
+            )
+            model_output = gr.Textbox(
+                label="Model Output", value="Please wait...", lines=4
+            )
+            # with gr.Column():
+            #     clazz = gr.Label()
+            # interpret_button = gr.Button("Interpret (very slow)")
+            # interpretation = gr.components.Interpretation(disassembly)
 
     def file_change_fn(file):
 
         if file is None:
-            return {
-                    col: gr.update(visible=False),
-                    all_dis_state: None
-                    }
+            return {col: gr.update(visible=False), state: {"file": None}}
         else:
 
             try:
                 progress = gr.Progress()
-                progress(0, desc="Analyzing binary...")
+                progress(
+                    0,
+                    desc=f"Analyzing binary  {os.path.basename(file.name)} with Ghidra...",
+                )
                 fun_data = get_functions(file.name)
-                #print(fun_data)
+                # print(fun_data)
 
-                addrs = [(f"{name} ({hex(int(addr))})", int(addr)) for addr, name in fun_data.items()]
-            except:
-                raise gr.Error("Unable to obtain functions")
+                addrs = [
+                    (f"{name} ({hex(int(addr))})", int(addr))
+                    for addr, name in fun_data.items()
+                ]
+            except Exception as e:
+                raise gr.Error(f"Unable to analyze binary with Ghidra: {e}")
 
             return {
-                    col: gr.Column(visible=True),
-                    fun_dropdown: gr.Dropdown(choices=addrs, value=addrs[0][1]),
-                    all_dis_state: fun_data
-                    }
+                col: gr.Column(visible=True),
+                fun_dropdown: gr.Dropdown(choices=addrs, value=addrs[0][1]),
+                state: {"file": file},
+            }
+
+    def function_change_fn(selected_fun, state, progress=gr.Progress()):
+
+        # disassembly_str = fun_data[int(selected_fun, 16)].decode("utf-8")
+        # load_results = model.fn(disassembly_str)
+        # top_k = {e['label']: e['confidence'] for e in load_results['confidences']}
+
+        with tempfile.TemporaryDirectory() as TEMP_DIR:
+
+            print(selected_fun)
+
+            progress(0, desc=f"Running DIRTY Ghidra on {hex(selected_fun)}...")
+
+            try:
+                subprocess.run(
+                    f"/ghidra/support/analyzeHeadless {TEMP_DIR} Project -import {state['file'].name} -postscript /DIRTY/scripts/DIRTY_infer.py {TEMP_DIR}/funcs.json {selected_fun}",
+                    shell=True,
+                )
+
+                json_info = json.load(open(f"{TEMP_DIR}/funcs.json"))
+
+                if "exception" in json_info:
+                    raise gr.Error(f"DIRTY Ghidra failed: {json_info['exception']}")
+
+            except Exception as e:
+                raise gr.Error(f"Unable to run DIRTY Ghidra: {e}")
+
+            #print(json_info)
+
+            return {
+                disassembly: gr.Textbox(value=json_info["disassembly"]),
+                original_decompile: gr.Textbox(value=json_info["original_decompile"]),
+                decompile: gr.Textbox(value=json_info["decompile"]),
+                model_output: gr.Textbox(value=json_info["model_output"]),
+            }
 
     # Need to put intro as output to get progress to work!
-    file_widget.change(file_change_fn, file_widget, outputs=[intro, col, fun_dropdown, all_dis_state])
+    file_widget.change(
+        file_change_fn, file_widget, outputs=[intro, state, col, fun_dropdown]
+    )
+
+    fun_dropdown.change(
+        function_change_fn,
+        inputs=[fun_dropdown, state],
+        outputs=[disassembly, original_decompile, decompile, model_output],
+    )
 
 # spaces only shows stderr..
 os.dup2(sys.stdout.fileno(), sys.stderr.fileno())