Filter out trivial examples

main.py

@@ -1,5 +1,4 @@
 import gradio as gr
-import shutil
 import subprocess
 import tempfile
 import os
@@ -107,15 +106,17 @@ with gr.Blocks() as demo:
 
                 addrs = [
                     (f"{name} ({hex(int(addr))})", int(addr))
-                    for addr, name in fun_data.items()
+                    for addr, (name, cf) in fun_data.items()
                 ]
+
+                cfs = {name: cf for (name, cf) in fun_data.values()}
             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]),
-                state: {"file": file},
+                state: {"file": file, "cfs": cfs},
             }
 
     def function_change_fn(selected_fun, state, progress=gr.Progress()):
@@ -131,18 +132,20 @@ with gr.Blocks() as demo:
             progress(0, desc=f"Running DIRTY Ghidra on {hex(selected_fun)}...")
 
             try:
-                subprocess.run(
+                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}",
                     shell=True,
                 )
+            except Exception as e:
+                raise gr.Error(f"Unable to run Ghidra: {e}\n{output}")
 
-                json_info = json.load(open(f"{TEMP_DIR}/funcs.json"))
+            if not os.path.exists(f"{TEMP_DIR}/funcs.json"):
+                raise gr.Error(f"DIRTY Ghidra failed to produce output: {output}")
 
-                if "exception" in json_info:
-                    raise gr.Error(f"DIRTY Ghidra failed: {json_info['exception']}")
+            json_info = json.load(open(f"{TEMP_DIR}/funcs.json"))
 
-            except Exception as e:
-                raise gr.Error(f"Unable to run DIRTY Ghidra: {e}")
+            if "exception" in json_info:
+                raise gr.Error(f"DIRTY Ghidra failed: {json_info['exception']}")
 
             #print(json_info)
 

scripts/dump_functions.py

@@ -1,4 +1,11 @@
 import json
+import sys
+
+sys.path.append("/DIRTY/dirty")
+
+import utils.infer
+from utils.dataset import Example
+
 
 def dump_functions_to_json():
     functionManager = currentProgram().getFunctionManager()
@@ -15,8 +22,15 @@ def dump_functions_to_json():
         func_name = func.getName()
         func_address = func.getEntryPoint().getOffset()
         
+        cf = utils.infer.ghidra_obtain_cf(func)
+
+        example = Example.from_cf(
+            cf, binary_file="binary_file", max_stack_length=1024, max_type_size=1024
+        )
+
         # Add function name and address to the dictionary
-        functions_dict[func_address] = func_name
+        if example.is_valid_example:
+            functions_dict[func_address] = (func_name, cf.to_json())
     
     # Convert the dictionary to a JSON object
     json_data = json.dumps(functions_dict, indent=4)