ejschwartz/function-method-detector

1import gradio as gr2 3import os4import re5import subprocess6import tempfile7from transformers import pipeline8 9 10MODEL_ID = "ejschwartz/oo-method-test-model-bylibrary"11 12classifier = pipeline(13    "text-classification",14    model=MODEL_ID,15)16 17def run_model(text):18    results = classifier(text, top_k=None, truncation=True)19    if isinstance(results, dict):20        results = [results]21    if results and isinstance(results[0], list):22        results = results[0]23 24    confidences = [25        {"label": entry["label"], "confidence": entry["score"]}26        for entry in results27    ]28    best_label = max(confidences, key=lambda entry: entry["confidence"])["label"] if confidences else "unknown"29    return {"label": best_label, "confidences": confidences}30 31def get_all_dis(bname, addrs=None):32 33    anafile = tempfile.NamedTemporaryFile(prefix=os.path.basename(bname) + "_", suffix=".bat_ana")34    ananame = anafile.name35 36    addrstr = ""37    if addrs is not None:38        addrstr = " ".join([f"--function-at {x}" for x in addrs])39 40    subprocess.check_output(f"bat-ana {addrstr} --no-post-analysis -o {ananame} {bname} 2>/dev/null", shell=True)41 42 43    output = subprocess.check_output(f"bat-dis --no-insn-address --no-bb-cfg-arrows --color=off {ananame} 2>/dev/null", shell=True)44    output = re.sub(b' +', b' ', output)45 46    func_dis = {}47    last_func = None48    current_output = []49 50    for l in output.splitlines():51        if l.startswith(b";;; function 0x"):52            if last_func is not None:53                func_dis[last_func] = b"\n".join(current_output)54            last_func = int(l.split()[2], 16)55            current_output.clear()56 57        if not b";;" in l:58            current_output.append(l)59 60    if last_func is not None:61        if last_func in func_dis:62            print("Warning: Ignoring multiple functions at the same address")63        else:64            func_dis[last_func] = b"\n".join(current_output)65 66    return func_dis67 68def get_funs(f):69    funs = get_all_dis(f.name)70    return "\n".join(("%#x" % addr) for addr in funs.keys())71 72with gr.Blocks() as demo:73 74    all_dis_state = gr.State()75 76    gr.Markdown(77        """78    # Function/Method Detector79 80    First, upload a binary.81 82    This model was only trained on 32-bit MSVC++ binaries.  You can provide83    other types of binaries, but the result will probably be gibberish.84    """85    )86 87    file_widget = gr.File(label="Binary file")88 89    with gr.Column(visible=False) as col:90        #output = gr.Textbox("Output")91 92        gr.Markdown("""93        Great, you selected an executable!  Now pick the function you would like to analyze.                    94                    """)95 96        fun_dropdown = gr.Dropdown(label="Select a function", choices=["Woohoo!"], interactive=True)97 98        gr.Markdown("""99        Below you can find the selected function's disassembly, and the model's100        prediction of whether the function is an object-oriented method or a101        regular function.102                    """)103 104        with gr.Row(visible=True) as result:105            disassembly = gr.Textbox(label="Disassembly", lines=20)106            with gr.Column():107                clazz = gr.Label()108 109    example_widget = gr.Examples(110        examples=[f.path for f in os.scandir(os.path.join(os.path.dirname(__file__), "examples"))],111        inputs=file_widget,112        outputs=[all_dis_state, disassembly, clazz]113    )114 115    def file_change_fn(file, progress=gr.Progress()):116 117        if file is None:118            return {col: gr.update(visible=False),119                    all_dis_state: None}120        else:121 122            #fun_data = {42: 2, 43: 3}123            progress(0, desc="Disassembling executable")124            fun_data = get_all_dis(file.name)125 126            addrs = ["%#x" % addr for addr in fun_data.keys()]127            default_addr = addrs[0] if addrs else None128 129            return {col: gr.update(visible=True),130                    fun_dropdown: gr.update(choices=addrs, value=default_addr),131                    all_dis_state: fun_data132                    }133        134    def function_change_fn(selected_fun, fun_data):135 136        disassembly_str = fun_data[int(selected_fun, 16)].decode("utf-8")137 138        load_results = run_model(disassembly_str)139        top_k = {e['label']: e['confidence'] for e in load_results['confidences']}140 141        return {disassembly: gr.update(value=disassembly_str),142                clazz: gr.update(value=top_k),143        }144 145    file_widget.change(file_change_fn, file_widget, [col, fun_dropdown, all_dis_state])146 147    fun_dropdown.change(function_change_fn, [fun_dropdown, all_dis_state], [disassembly, clazz])148 149demo.queue()150demo.launch(151    server_name="0.0.0.0",152    server_port=7860,153    #share=True,154    debug=True,155    show_error=True,156)157